From 629a41acca6ebc80f78c681c8aa1558f16f0866b Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 23:18:04 +0200 Subject: [PATCH 01/39] =?UTF-8?q?feat(data):=20one=20Resample=20=E2=80=94?= =?UTF-8?q?=20a=20stored=20transform=20decoded=20and=20streamed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A geometry vocabulary (grids, boxes, affine maps); a decoder that splits a stored ITK transform into an affine stage and a bounded displacement residual; a torch sampler that applies it on the volume's device. Resample becomes one transform: a target grid — its own, derived from the source by ITK's own arithmetic, or a reference case — reached through the declared map, streamed slab by slab, and it says why when it cannot stream. A map that factorises is read one axis at a time and blended most-shrinking axis first; one that does not goes through grid_sample. A B-spline order with no kernel where the value is built is refused. Docs say what the old names now mean, and the doc harness runs against a real stored transform. --- AGENTS.md | 2 +- CHANGELOG.md | 62 + docs/source/concepts/streaming.md | 16 +- docs/source/config_guide/transform.md | 142 +- docs/source/reference/api/extension-points.md | 1 - .../source/reference/components/transforms.md | 11 +- konfai/data/data_manager.py | 10 +- konfai/data/geometry.py | 462 ++++ konfai/data/patching.py | 71 +- konfai/data/sampling.py | 454 ++++ konfai/data/transform.py | 1885 +++++++---------- konfai/predictor.py | 75 +- konfai/utils/ITK.py | 118 ++ .../test_transform_doc_examples.py | 18 + tests/unit/test_geometry.py | 178 ++ tests/unit/test_resample.py | 383 ++++ tests/unit/test_resample_sampler_rules.py | 137 +- tests/unit/test_resample_to_reference.py | 181 +- tests/unit/test_resample_transform.py | 340 +++ tests/unit/test_sampling.py | 280 +++ tests/unit/test_streamed_read_dispatcher.py | 21 +- tests/unit/test_streamed_write_dispatcher.py | 52 +- tests/unit/test_transform.py | 33 +- tests/unit/test_transform_bound.py | 255 +++ .../unit/test_transform_locality_contract.py | 50 +- tests/unit/test_transformer_workflow.py | 14 +- tests/unit/test_warp.py | 73 +- .../test_write_pyramid_and_field_bound.py | 44 +- 28 files changed, 3927 insertions(+), 1441 deletions(-) create mode 100644 konfai/data/geometry.py create mode 100644 konfai/data/sampling.py create mode 100644 tests/unit/test_geometry.py create mode 100644 tests/unit/test_resample.py create mode 100644 tests/unit/test_resample_transform.py create mode 100644 tests/unit/test_sampling.py create mode 100644 tests/unit/test_transform_bound.py diff --git a/AGENTS.md b/AGENTS.md index d1ab1009..27e3ef2d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -61,7 +61,7 @@ Every extension point is **"subclass a base, reference it by classpath in YAML"* - **Model:** subclass `network.Network`, build the graph in `__init__` via `add_module`. Reference `classpath: module.MyNet`, a local `Model:MyNet`, a `.yml`, or `default|.yml` for the shipped catalog. - **Pretrained weights:** `utils/pretrained.py:transfer_weights_by_execution_order` pairs weighted leaves in forward-execution order (no key map). It fills **every** target tensor or raises — a tensor held by a parent module (`torch.nn.MultiheadAttention` owns `in_proj_weight` beside its `out_proj` child) or by a submodule the forward skips cannot be paired. Unreached *source* branches (nnU-Net deep-supervision heads) are ignored on purpose. - **Loss / metric:** subclass `metric.measure.Criterion`; `forward` returns a `Tensor` (loss) or a `(value, dict)` tuple (metric — consumers `isinstance`-branch). Attach under `outputs_criterions`/`metrics` to a **named module output**. Optional-dep criteria import lazily via `_require_optional(...)` and raise an actionable `MeasureError` — never a bare top-level import. -- **Transform:** subclass `data.transform.Transform`; implement `__call__` **and** `transform_shape()` (must predict the output spatial shape *exactly* — patch planning depends on it). Declare `patch_locality()` (a `LocalityKind`: `POINTWISE`/`HALO`/`ORIENTATION`/`CROP`/`GLOBAL_STAT`/`RESCALE`/`SLAB`/`WHOLE_VOLUME`) or the base default makes it `WHOLE_VOLUME`; a `WHOLE_VOLUME` that is a property of the *configuration* rather than of the stage must carry `reason=`, which the plan prints — without it the reader has nothing to change. Pair `inverse()` if `apply_inverse`; override `prepare(konfai_args)` only when the stage builds a sub-object from configuration of its own (`Reduce` → its operator). +- **Transform:** subclass `data.transform.Transform`; implement `__call__` **and** `transform_shape()` (must predict the output spatial shape *exactly* — patch planning depends on it). Declare `patch_locality()` (a `LocalityKind`: `POINTWISE`/`HALO`/`ORIENTATION`/`CROP`/`GLOBAL_STAT`/`REGRID`/`SLAB`/`WHOLE_VOLUME`) or the base default makes it `WHOLE_VOLUME`; a `WHOLE_VOLUME` that is a property of the *configuration* rather than of the stage must carry `reason=`, which the plan prints — without it the reader has nothing to change. Pair `inverse()` if `apply_inverse`; override `prepare(konfai_args)` only when the stage builds a sub-object from configuration of its own (`Reduce` → its operator). - **Augmentation:** subclass `data.augmentation.DataAugmentation`; `_state_init` (sample params per case index) + `_compute` (apply lazily). Only `Mask`/`Permute` may change shape. A draw is also a **chain stage**: `TransformLoader` resolves a bare name against `data.transform` first and `data.augmentation` second, so a `transforms:` block may interleave draws and transforms — which is how TRANSFORM declares per-copy draws after an `Expand`. - **Reduction:** subclass `data.reduction.Reduction`; implement `__call__(list[Tensor]) -> Tensor` over the `[1, K, C, *spatial]` layout both engines hand over. Two consumers, one vocabulary: the predictor folds one case's copies (ensemble/TTA), `data.transform.Reduce` folds N **cases** into one. Declare `voxel_local = True` only if every output voxel reads the same voxel of each input (**a wrong `True` corrupts a streamed output** — the gate checks nothing else), `incremental = True` if `accumulate` can fold one at a time, and override `output_channels(channels, cases)` when the fold changes the channel count (`Concat` does). `Reduce` refuses a non-`voxel_local` operator outright. - **Imaging format:** add a `Dataset.AbstractFile` backend, dispatch it in `File.__enter__`, register aliases in `SUPPORTED_EXTENSIONS`; import-guard the heavy lib. diff --git a/CHANGELOG.md b/CHANGELOG.md index fa86fbde..8e2bad0b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,68 @@ draft, then say what a user of the package gets that they did not have -- and re against the commits that landed *after* you drafted it. Running the command over a section already written replaces it. +## Unreleased + +### ✨ Features + +- **transform**: one `Resample`. `ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, + `ResampleTransform` and `Warp` were five stages answering two questions between them, each with a + sampler of its own. They are now five spellings of one stage that asks the questions separately: + **which grid to write on** (nothing, `spacing`, `shape`, or `reference` — a stored image's grid + adopted whole) and **what map to write it through** (`field`, `transforms`, or neither). Every + combination is legal, and asked for together they compose into **one interpolation** instead of + two. The old names still work and are thin argument translations. +- **transform**: `align` says where a `spacing` or `shape` grid sits — `extent` (the default) keeps + the field of view, `origin` keeps voxel zero's centre. This was decided silently before, and + differently by the data and by the header. +- **transform**: a resample streams whatever it is asked for, and on the volume's device. Applying a + stored registration used to hold a whole volume — not because a warp needs one, but because + nothing on a `sitk.Transform` said how far it reached. A rigid or affine map now bounds exactly; a + BSpline and a displacement field bound by the largest of their values, which holds at every point + rather than at the sampled ones. Everything is evaluated in torch, so nothing marshals a + GPU-resident case out to numpy and back. +- **transform**: a resample no longer requires the grids to share a direction. A rotated reference, + or a field stored on turned axes, used to be refused with an instruction to run `Canonical` first + — a second interpolation of the same voxels. Both now work directly. +- **transform**: a field is read on its own grid. `Warp` required the field and the case to share + one; a field solved at 120 µm now moves a volume stored at 30 µm without being upsampled first. +- **transform**: `ResampleToShape` needs no geometry at all. A count is a count; only a change of + density needs the density it starts from. + +### 🐛 Fixes + +- **transform**: a resampled label map no longer comes out shifted against the image beside it. + `F.interpolate`'s nearest reads `floor(o * scale)` where its linear reads `scale * (o + 0.5) - 0.5`, + so a mask resampled by the same stage as its CT lagged it by `(scale - 1) / 2` source voxels — + 2.5 voxels, 1.25 mm of anatomy, resampling 0.5 mm to 3 mm. Both volumes were entirely plausible on + their own. Nearest is now ITK's round-half-up on the same physical index the linear sampler reads. +- **transform**: the header a resample records now describes the grid it actually sampled. + `ResampleToResolution` wrote the spacing that was *asked for* while sampling at `n_in/n_out` times + the source's (up to a millimetre of drift across a volume) and left the `Origin` alone while + sampling half a spacing-change away from it. Nothing downstream could see either: the voxels are + all real, and the header was the only witness. +- **transform**: a voxel count no longer loses a slice to floating point. 90 voxels of 0.7 mm re-cut + at 1.5 mm is 42, and the count went through float32 to get there — landing on 41 or 42 depending + on the numbers. +- **transform**: a warp on an oblique case reads the neighbourhood it needs. The halo was derived + per array axis from a world displacement, which assumes the direction cosines are the identity; on + a turned case the window was short on the axes the displacement actually reached, and a short + window returns the border value rather than raising. +- **transform**: a resample refuses what it used to do quietly — resampling in a physical space the + case does not have, applying a transform type nothing bounds, inverting a spline or a field by + building a whole-grid displacement field and iterating on it per case, and warping through a field + with no declared bound. Each declares `WHOLE_VOLUME` with the sentence saying what to change, and + the run proceeds on the whole-volume path. +- **transform**: `ResampleTransform`'s `inverse` defaults to `false`. It always raised + `NotImplementedError`, so a prediction finalize through this stage failed at the end of the run + rather than at its configuration. + +### 🔧 Internals + +- **data**: `LocalityKind.RESCALE` is gone. It was the dispatcher's own resample map — a size ratio, + which says nothing once a target grid has an origin — and with one resample stage there is one + regime, `REGRID`, that the stage owns both halves of. + ## v1.8.0 (2026-08-04) ### ✨ Features diff --git a/docs/source/concepts/streaming.md b/docs/source/concepts/streaming.md index 82ae1c8c..23a275f1 100644 --- a/docs/source/concepts/streaming.md +++ b/docs/source/concepts/streaming.md @@ -99,7 +99,6 @@ which region of the file a patch needs. | `ORIENTATION` | flip or permute | the index-remapped region | | `CROP` | the source region is the target translated | the region — reading it *is* the answer | | `GLOBAL_STAT` | needs whole-volume `Min`/`Max`/`Mean`/`Std` | the statistic once from disk, then the exact patch | -| `RESCALE` | resample | the region through the scale mapping, plus an interpolation halo | | `SLAB` | a per-voxel value map plus a side effect that needs the slab's place in the volume | nothing on the read path — the dispatcher treats it as `WHOLE_VOLUME`; it streams on the write side (`Mask`, `InferenceStack`) | | `WHOLE_VOLUME` | genuinely needs everything | the volume — the fallback | @@ -107,7 +106,7 @@ which region of the file a patch needs. volume and is correct. A chain streams when every stage is pointwise or a region kind — `HALO`, -`ORIENTATION`, `CROP`, `RESCALE` — where `GLOBAL_STAT` counts as pointwise. +`ORIENTATION`, `CROP`, `REGRID` — where `GLOBAL_STAT` counts as pointwise. Region stages **compose** in any number: each stage's region pulls through the one before it, down to one bounded read of the stored volume. The chain planned is the group's `transforms` followed by the copy's own augmentation draw — one @@ -121,9 +120,8 @@ Seven conditions reject streaming: 2. a halo wider than half the read extent on any axis; 3. a `GLOBAL_STAT` preceded by a stage that does not preserve statistics; 4. a `GLOBAL_STAT` whose statistic cannot be read from disk; -5. a `RESCALE` declared by a stage that is not a `Resample` subclass; -6. a `RESCALE` on a case with no `Spacing`; -7. a chain whose folded shapes do not land on the target grid. +5. a `REGRID` whose stage cannot size the region it reads (no geometry, no bound); +6. a chain whose folded shapes do not land on the target grid. ### Why `[Clip(-200, 400), Standardize()]` does not stream @@ -185,7 +183,7 @@ streams and `Dilate(5)` does not. | `HALO` | `Dilate(n>0)`, `Gradient` | | `ORIENTATION` | `Flip`, `Permute`, `Canonical` (only on axis-aligned direction cosines) | | `CROP` | `Crop` (only once its box is on the case) | -| `RESCALE` | `ResampleToResolution`, `ResampleToShape` | +| `REGRID` | `Resample` — and its spellings `ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, `ResampleTransform`, `Warp` | Augmentations declare per **(case, draw)** — two copies of the same case can answer differently. `Permute`, `Flip` (when `vector_field: false`), and `Rotate` @@ -270,7 +268,7 @@ To opt in, override `patch_locality` under three rules: `ORIENTATION` and `CROP` must also implement `stream_region_source`, mapping a target patch to the source region. Declaring a region kind without it raises a -`TransformError` on the first patch read. `HALO` and `RESCALE` need no remap; the +`TransformError` on the first patch read. `HALO` needs no remap; the dispatcher derives their regions. ## Equivalence @@ -287,12 +285,12 @@ land a few float32 ulp away. A stage seeded from `Min`/`Max` — `Normalize`, `Clip('min', 'max')` — is byte-identical: a min and a max have no summation order to disagree on. -A streamed `RESCALE` computes interpolation weights in the sub-region's +A streamed `REGRID` computes interpolation weights in the sub-region's coordinate frame, so the deviation scales with the local gradient rather than with the voxel's own value: within a few ulp of the volume's peak on float volumes, within 1 LSB on integer ones. A nearest-neighbour resample — what a `uint8` label volume gets — uses no weights and stays exact. The `Translate` -augmentation interpolates on the same terms as `RESCALE`. +augmentation interpolates on the same terms as `REGRID`. ## Next steps diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index e889f859..d22da3fe 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -19,7 +19,7 @@ Transformer: groups_dest: CT_iso: transforms: - ResampleToResolution: + Resample: spacing: [1.0, 1.0, 1.0] Write: dataset: ./Out:omezarr @@ -95,7 +95,7 @@ few percent of the budget can still exceed it. | --- | --- | | `allow` | Take the whole-volume path silently — but the plan still names it. | | `warn` (default) | Same, plus a warning line after the plan. | -| `error` | Refuse the run. Nothing is written. A fallback only discovered mid-run (a failed sweep, a `Warp` bound exceeded) stops at that case — earlier cases stay written, and the per-case resume covers the rerun. | +| `error` | Refuse the run. Nothing is written. A fallback only discovered mid-run (a failed sweep, a `Resample` field bound exceeded) stops at that case — earlier cases stay written, and the per-case resume covers the rerun. | Independently of `on_fallback`, a case that **cannot stream and does not fit `memory_budget`** always refuses the whole run, before the first byte. Writing @@ -258,7 +258,7 @@ case's chain *lands* on (a `Resample` before the `Reduce` counts): Nothing can verify that the members truly live in a common space — only that they claim to. `shape_only` and `reference:` will happily average misaligned volumes; the result still looks like a volume and is an artefact. Put the -cohort on one grid first — `ResampleToReference`, below. +cohort on one grid first — `Resample: {reference: …}`, below. ``` **A reduction has no whole-volume fallback.** Folding every case in memory is @@ -268,33 +268,66 @@ voxel-local stages (and statistics, seeded by an extra pass of the engine's own) may follow the `Reduce` in the same chain: anything reading across space belongs in a second chain that reads the written output back. -### `ResampleToReference`: making `strict` true rather than waived +### `Resample`: one stage, two questions + +Every resample in KonfAI is one stage, `Resample`, and it asks two independent +questions: + +| | key | meaning | +| --- | --- | --- | +| **which grid to write on** | *(nothing)* | the case's own — the map moves the anatomy, the voxels stay put | +| | `spacing` | the same field of view at another density | +| | `shape` | the same field of view at a given count | +| | `reference` | the grid of a stored image, adopted whole | +| **what map to write it through** | *(nothing)* | the identity — a change of grid and nothing else | +| | `field` | a displacement field, on its own grid, in world units | +| | `transforms` | transforms stored beside the cases (rigid, affine, BSpline, field, composite) | + +Any combination is legal, and asked for together they compose into **one +interpolation**: the source is read once, at the displaced point. Doing it as +two stages resamples twice, and a volume interpolated twice has lost detail the +second pass invented no more of. + +`align` places a `spacing` or `shape` grid, and it is the one choice the family +used to make silently: `extent` (the default) keeps the field of view — the +outer faces coincide — while `origin` keeps voxel zero's centre where it is. A +quarter of a voxel of anatomy separates them, and a `reference` states its own +placement and ignores this. + +```{note} +`ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, +`ResampleTransform` and `Warp` are still accepted, and are now thin spellings of +this one stage: `Resample: {spacing: …}`, `{shape: …}`, `{reference: …}`, +`{transforms: …}` and `{field: …}` respectively. +``` + +### `Resample: {reference: …}`: making `strict` true rather than waived 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 +coordinates are not an anatomical frame. A `reference` grid 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} + Resample: {reference: case_0, reference_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 +The reference is a stored image, named by `reference` — and by +`reference_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} + Resample: {reference: case_1, reference_group: CT, reference_dataset: ./Raw:mha} Write: {dataset: ./OnTemplate:mha} ``` -`dataset:` takes the same `path[:format]` spec as everywhere else, so the +`reference_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. @@ -307,9 +340,9 @@ position, added, and the source sampled **once** at the displaced point: ```yaml transforms: - ResampleToReference: - entry: case_0 - group: CT + Resample: + reference: case_0 + reference_group: CT field: ./Fields:mha field_group: DVF max_displacement: 4.0 @@ -319,7 +352,7 @@ transforms: 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 +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 @@ -337,13 +370,22 @@ identity where the field says nothing, as SimpleITK has it. 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. +declares `WHOLE_VOLUME` and says so in the plan. + +Naming no target grid is the shape update of an atlas build — the field applied +on the case's *own* grid — and is the same stage with `reference` left out: + +```yaml +transforms: + Resample: {field: ./Fields:mha, field_group: DVF, max_displacement: 4.0} + Write: {dataset: ./Warped:mha} +``` ```{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. +This was `Warp`, which required the field and the case to share a grid. They no +longer have to: the field is read at each target voxel's world position on the +field's own grid, so a field solved at 120 µm moves a volume stored at 30 µm +without being upsampled first. ``` Naming an image rather than fifteen numbers is deliberate. A grid is an extent @@ -365,7 +407,7 @@ on its own — a CT is `int16` and so is nothing else about it — so a label ma stored as anything but `uint8` must say so: ```yaml -ResampleToReference: {entry: case_0, group: Labels, interpolation: nearest} +Resample: {reference: case_0, reference_group: Labels, interpolation: nearest} ``` Getting it wrong is silent. Two labels blended give a third that was never in @@ -386,6 +428,59 @@ 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. +### `Resample: {transforms: …}`: applying a registration that was already solved + +With no target grid named, `Resample` changes nothing about the grid and moves +the anatomy through transforms **stored beside the cases** — the apply step of a +registration solved elsewhere: + +```yaml +transforms: + Resample: {transforms: {reg: false}} + Write: {dataset: ./Registered:mha} +``` + +Each key of `transforms:` is a group of the run's own datasets holding one +transform per case; the value says whether to invert it. Rigid, affine, BSpline +and displacement-field entries are all read the same way, and several groups +compose — the **last declared is applied first**, which is SimpleITK's own +composite order. + +**It streams**, and what makes that possible is that a stored transform can say +how far it reaches. A rigid or affine map is an exact affine, so the source box +of a target region is that region's box mapped through it. A BSpline and a dense +field are values on a grid read through a kernel that is non-negative and sums +to one, so the largest of those values bounds the displacement at *every* point +— not at the points someone sampled. The region a slab must read is therefore +known before a voxel is touched. + +```{warning} +Bounded is not the same as cheap. A map **oblique to the storage axes** has an +axis-aligned source box that covers most of the volume on two axes, and it gets +worse the thinner the slabs: the same case that reads 1.0× its bytes in one +piece reads several times that in slabs. The bound is exact either way — this is +a property of the decomposition, not a defect — but it is why streaming such a +map is not automatically worth it. Bring the case onto an axis-aligned grid +first (`Canonical`) when the geometry allows it. +``` + +**What it refuses**, declaring `WHOLE_VOLUME` with the reason so the run +proceeds on the whole-volume path rather than breaking: + +- a case whose header carries no `Origin` / `Spacing` / `Direction` — a stored + transform is applied in physical space, and without a geometry there is none; +- a transform type that decomposes into no bounded map, naming the type; +- `invert: true` on a spline or a displacement field. Inverting one is a dense + solve over the whole grid, and a field solved per region is not the + restriction of the field solved once — so store the inverse where the + transform is written, or set the group to `false`. + +`interpolation` and `fill` work the same wherever they appear: unset, `uint8` +takes the nearest voxel and everything else is interpolated, and a label map +stored as anything else must say `interpolation: nearest`. Nearest here is ITK's +round-half-up on the physical index — the same coordinate the linear sampler +reads, so a mask and the image beside it land on the same voxels. + ### `Expand`: one case, N copies `Expand` multiplies, and nothing else. The draws are **ordinary stages of the @@ -413,7 +508,7 @@ Transformer: pattern: "{name}_r{a:02d}" Rotate: # a draw, per copy is_quarter: true - ResampleToResolution: # a transform, per copy + Resample: # a transform, per copy spacing: [2.0, 2.0, 2.0] Brightness: # another draw, per copy b_std: 0.2 @@ -510,7 +605,7 @@ picks a regime per copy and the plan prints which: - **WHOLE-VOLUME** — the copy's chain cannot stream at all; the shared part is still assembled only once for the case. -When the shared prefix is expensive (a `Warp`, a resample), put a `Save` before +When the shared prefix is expensive (a resample, a warp), put a `Save` before the `Expand`: it is materialized once and every copy reads the cache. Resume is **per copy**: a copy whose entry exists is skipped, so an interrupted @@ -613,8 +708,7 @@ 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()` | -| the same box, resampled | `RESCALE` | inherit from `Resample` | -| another grid entirely | `REGRID` | `stream_region_source()` and `stream_region()` | +| another grid, or the same one at another density | `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/docs/source/reference/api/extension-points.md b/docs/source/reference/api/extension-points.md index d1891b23..56104206 100644 --- a/docs/source/reference/api/extension-points.md +++ b/docs/source/reference/api/extension-points.md @@ -125,7 +125,6 @@ halo of a geometric draw is that draw's own. | `ORIENTATION` | flip or permute | `stream_region_source` | | `CROP` | source region is the target region translated | `stream_region_source` | | `GLOBAL_STAT` | needs whole-volume statistics, `stat_keys` a subset of Min/Max/Mean/Std (or their `…PerChannel` forms) | nothing — the dispatcher seeds the statistic from disk | -| `RESCALE` | resample by a ratio | subclass `Resample` | | `REGRID` | resample onto a grid declared elsewhere — a stored reference, not a ratio — so the source region is computed from the two geometries | subclass `Resample`; declare a halo when a displacement field is composed in | | `SLAB` | a per-voxel value map plus a side effect that needs the slab's place in the volume | `stream_slab(name, tensor, region, spatial_shape, cache_attribute)`, and optionally `stream_abort`. The **read** dispatcher has no slab context and treats it as `WHOLE_VOLUME`; the gain is on the write side | | `WHOLE_VOLUME` | needs the whole volume | nothing — this is the default | diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 4c3cd1d9..6a9fe829 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -97,11 +97,12 @@ until it declares otherwise. | --- | --- | --- | --- | --- | --- | | `Padding` | `F.pad`; updates Origin. `mode` supports `"constant:"`. | `padding=[0,0,0,0,0,0], mode="constant", inverse=True` | **yes** | **yes** | no‡ | | `Crop` | Crop to foreground bounding box; caches the box; updates Origin. | `inverse=True` | **yes** | **yes** (pads back) | **yes** — once the `box` is on the case; the region is the patch translated | -| `ResampleToResolution` | Resample to a target voxel spacing (per-axis `<0` = keep). | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** — resampled from the source region | -| `ResampleToShape` | Resample to a target shape (per-axis `0/<0` = keep). | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** — resampled from the source region | -| `ResampleTransform` | Warp by stored SimpleITK transforms read from the dataset. | `transforms`, `inverse=True` | no | no | no — nothing bounds how far the stored displacement reaches | -| `ResampleToReference` | Resample onto the grid of a **declared reference case** — extent, spacing, origin and direction — so a cohort meets on a grid that is one of its own rather than an invented one. An optional `field` composes the grid change and the warp into one pass, so the intermediate never exists. `interpolation` left unset is nearest for `uint8` and linear otherwise; declare it for a label map stored as anything else. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** — the grid change alone; a composed field is not inverted | **yes** — declares `REGRID`; halo = bound / spacing when a field is composed in | -| `Warp` | Resample a case through a displacement field on the **same grid** — the shape update of an atlas build. Warping onto a different grid is a resample too, and is not this stage. The declared bound is CHECKED per component against every region read, so a field that exceeds it raises instead of sampling zeros (which would read as a dark rim and nothing else). | `field` (required), `group=None`, `max_displacement=0.0`, `interpolation="linear"` | no | no | **yes** — halo = bound / spacing. `max_displacement: auto` reads the bound the fields recorded when KonfAI wrote them (OME-Zarr only); with no bound at all it declares whole-volume and says which one is missing | +| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk bounds by `max_displacement`, **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, `invert: true` names a spline or a field, or a field has no bound | +| `ResampleToResolution` | Deprecated spelling of `Resample: {spacing: ...}`. | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** | +| `ResampleToShape` | Deprecated spelling of `Resample: {shape: ...}`. | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** | +| `ResampleToReference` | Deprecated spelling of `Resample: {reference: ...}`. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** | **yes** | +| `ResampleTransform` | Deprecated spelling of `Resample: {transforms: ...}`. | `transforms` (required), `interpolation=None`, `fill=0.0`, `inverse=False` | no | no | **yes** | +| `Warp` | Deprecated spelling of `Resample: {field: ...}`. Note that `Resample` no longer requires the field and the case to share a grid. | `field` (required), `group=None`, `max_displacement=0.0`, `interpolation="linear"` | no | no | **yes** | | `Canonical` | Reorient to canonical direction (3-D); updates Origin/Direction. | `inverse=True` | **yes** — a remap that transposes extents moves the patch grid | **yes** | **yes** — when the case's direction is a signed axis permutation; no on an oblique one (it is resampled) | | `Permute` | Permute spatial axes. `dims` is a pipe-separated axis list. | `dims="1\|0\|2", inverse=True` | **yes** | **yes** | **yes** — index remap | | `Flip` | Flip spatial axes. | `dims="1\|0\|2", inverse=True` | no | **yes** (self-inverse) | **yes** — index remap | diff --git a/konfai/data/data_manager.py b/konfai/data/data_manager.py index 4d58d132..c6977286 100755 --- a/konfai/data/data_manager.py +++ b/konfai/data/data_manager.py @@ -189,14 +189,10 @@ def _check_patch_transform_locality(transform: Transform, group_src: str, group_ " crops that patch about its own extent, and cuts the patch grid predictions are" " reassembled onto down to what is left." ), - LocalityKind.RESCALE: ( - 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." + f"'{name}' resamples its input onto another grid: applied to one patch it would rescale" + " that patch about its own extent, or hand back the whole target extent -- neither of" + " which is the patch grid predictions are reassembled onto." ), LocalityKind.WHOLE_VOLUME: f"'{name}' needs the whole volume.", } diff --git a/konfai/data/geometry.py b/konfai/data/geometry.py new file mode 100644 index 00000000..ffb0662a --- /dev/null +++ b/konfai/data/geometry.py @@ -0,0 +1,462 @@ +# 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 + +"""Grids, boxes and affine maps in world coordinates — the value vocabulary a resample shares. + +Two axis orders coexist in every KonfAI header, and every geometry bug this file's history records +is a confusion between them: array data is ``(Z, Y, X)``, physical geometry (``Origin``, +``Spacing``, ``Direction``) is ``(x, y, z)``. The types here carry the order in the field name — +``size_zyx``, ``origin_xyz`` — so a mixed expression reads as wrong at the call site instead of +resampling perfectly onto the wrong place. + +Everything is plain float64 numpy: no torch, no SimpleITK. A value built here crosses the +``mp.spawn`` pickle boundary as data, and the SimpleITK plumbing that produces it lives in +``konfai.utils.ITK`` behind its import guard. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +from konfai.utils.errors import TransformError + +if TYPE_CHECKING: + from konfai.utils.dataset import Attribute + +#: The geometry keys a grid is read from, in the order refusals name them. +_GEOMETRY_KEYS = ("Origin", "Spacing", "Direction") + + +def _as_float(values: object, rank: int, what: str, count: int) -> np.ndarray: + array = np.asarray(values, dtype=np.float64).ravel() + if array.size != count: + raise TransformError( + f"{what} does not describe a {rank}-dimensional grid: got {array.size} value(s), expected {count}." + ) + return array + + +@dataclass(frozen=True) +class AffineMap: + """``q_xyz = matrix @ p_xyz + translation_xyz``, in world coordinates.""" + + matrix: np.ndarray + translation: np.ndarray + + @staticmethod + def identity(rank: int) -> AffineMap: + return AffineMap(np.eye(rank), np.zeros(rank)) + + @property + def rank(self) -> int: + return int(self.translation.size) + + @property + def is_identity(self) -> bool: + return bool(np.array_equal(self.matrix, np.eye(self.rank)) and not self.translation.any()) + + def apply(self, points_xyz: np.ndarray) -> np.ndarray: + """Map points of shape ``(..., rank)``, accumulating exactly as ITK does. + + ``translation + Σ_j column_j · p_j`` with ``j`` ascending — the association of + ``TransformIndexToPhysicalPoint``. A matmul sums in whatever order BLAS picks, which is + one ULP away on an oblique grid, and one ULP of origin is the difference between a + streamed slab that is bit-identical to the whole volume and one that is merely close. + """ + points = np.asarray(points_xyz, dtype=np.float64) + out = np.broadcast_to(self.translation, points.shape).copy() + for j in range(self.rank): + out += points[..., j, np.newaxis] * self.matrix[:, j] + return out + + def then(self, outer: AffineMap) -> AffineMap: + """The composition ``outer(self(p))``.""" + return AffineMap(outer.matrix @ self.matrix, outer.matrix @ self.translation + outer.translation) + + def inverted(self) -> AffineMap: + """The inverse map, or a refusal when the matrix is singular. + + Refused rather than returned as garbage: a singular matrix here means a degenerate grid or + transform, and ``pinv`` would hand back a map that resamples plausibly from the wrong place. + """ + try: + inverse = np.linalg.inv(self.matrix) + except np.linalg.LinAlgError: + raise TransformError( + "This affine map is singular and cannot be inverted.", + "The grid or stored transform behind it collapses at least one axis; check its" + " Spacing/Direction (or the transform's matrix) for a zero row.", + ) from None + return AffineMap(inverse, -inverse @ self.translation) + + +@dataclass(frozen=True) +class WorldBox: + """An axis-aligned box in world coordinates, ``low_xyz`` to ``high_xyz`` inclusive.""" + + low_xyz: np.ndarray + high_xyz: np.ndarray + + def grown(self, radius_xyz: np.ndarray | float) -> WorldBox: + radius = np.broadcast_to(np.asarray(radius_xyz, dtype=np.float64), self.low_xyz.shape) + return WorldBox(self.low_xyz - radius, self.high_xyz + radius) + + def image_under(self, affine: AffineMap) -> WorldBox: + """The axis-aligned hull of this box's image under ``affine``. + + Centre and half-extents: the image of centre ``c`` is ``A c + b``, and the largest reach of + ``A h`` over the corners is ``|A| h`` because each corner coordinate is ``±h_k``. Equal to + the hull of the ``2^rank`` mapped corners — pinned by a test against that enumeration — in + O(rank²) and with no corner loop to get wrong. + """ + centre = (self.low_xyz + self.high_xyz) / 2.0 + half = (self.high_xyz - self.low_xyz) / 2.0 + mapped = affine.apply(centre) + reach = np.abs(affine.matrix) @ half + return WorldBox(mapped - reach, mapped + reach) + + +#: How close to a whole number a voxel count must be before it is taken to BE that number. A +#: spacing of 0.7 mm is not representable in binary, so 90 voxels of it re-cut at 1.5 mm come to +#: 41.999999999999997 and truncate to 41 -- one slice of anatomy dropped, and a spacing recorded +#: that no longer covers what was read. The band is far narrower than any density a header states. +_COUNT_TOLERANCE = 1e-6 + + +def _voxel_count(extent: int, spacing: float, wanted: float) -> int: + """How many voxels of ``wanted`` size cover ``extent`` voxels of ``spacing`` — truncated.""" + return int(np.floor(extent * spacing / wanted + _COUNT_TOLERANCE)) + + +@dataclass(frozen=True) +class Grid: + """A stored volume's sampling grid: extent in array order, geometry in physical order.""" + + size_zyx: tuple[int, ...] + origin_xyz: np.ndarray + spacing_xyz: np.ndarray + direction_xyz: np.ndarray + + @classmethod + def identity(cls, spatial_shape: list[int]) -> Grid: + """The grid of a volume with no geometry: unit spacing, origin zero, axes as stored. + + What a header carrying no ``Origin``/``Spacing``/``Direction`` means when the question is + only a change of extent. Under it a world coordinate IS an index, so a resample onto another + grid degenerates to the size ratio it always was -- which is how one engine serves a case + whose geometry is known and one whose is not, instead of a physical path and a ratio path + that have to be kept agreeing. + """ + rank = len(spatial_shape) + return cls(tuple(int(extent) for extent in spatial_shape), np.zeros(rank), np.ones(rank), np.eye(rank)) + + @classmethod + def from_header(cls, spatial_shape: list[int], attribute: Attribute, what: str) -> tuple[Grid, frozenset[str]]: + """The grid a header describes AND which of its keys it did not say — never a refusal. + + The identity stands in for what is absent, and the caller is told what it stood in for: what + a missing key costs depends on the question. An extent change needs none of them; a density + change needs the ``Spacing`` and nothing else; a reference grid or a stored map needs a real + physical space and so needs all three. Deciding that here, once, for every caller would + either refuse a resample that was answerable or answer one that was not. + """ + rank = len(spatial_shape) + identity = cls.identity(spatial_shape) + missing = frozenset(key for key in _GEOMETRY_KEYS if key not in attribute) + origin = ( + identity.origin_xyz + if "Origin" in missing + else _as_float(attribute.get_np_array("Origin"), rank, f"the Origin of {what}", rank) + ) + spacing = ( + identity.spacing_xyz + if "Spacing" in missing + else _as_float(attribute.get_np_array("Spacing"), rank, f"the Spacing of {what}", rank) + ) + direction = ( + identity.direction_xyz + if "Direction" in missing + else _as_float(attribute.get_np_array("Direction"), rank, f"the Direction of {what}", rank * rank).reshape( + rank, rank + ) + ) + if not np.all(spacing > 0.0): + raise TransformError( + f"The Spacing of {what} is {spacing.tolist()}.", + "A spacing is a physical extent per voxel and must be positive on every axis.", + ) + return cls(identity.size_zyx, origin, spacing, direction), missing + + @classmethod + def of(cls, spatial_shape: list[int], attribute: Attribute, what: str) -> Grid: + """The grid a header describes, or a refusal naming ``what`` and the missing key.""" + missing = [key for key in _GEOMETRY_KEYS if key not in attribute] + if missing: + raise TransformError( + f"The geometry of {what} is needed 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).", + ) + rank = len(spatial_shape) + origin = _as_float(attribute.get_np_array("Origin"), rank, f"the Origin of {what}", rank) + spacing = _as_float(attribute.get_np_array("Spacing"), rank, f"the Spacing of {what}", rank) + direction = _as_float(attribute.get_np_array("Direction"), rank, f"the Direction of {what}", rank * rank) + if not np.all(spacing > 0.0): + raise TransformError( + f"The Spacing of {what} is {spacing.tolist()}.", + "A spacing is a physical extent per voxel and must be positive on every axis.", + ) + return cls(tuple(int(extent) for extent in spatial_shape), origin, spacing, direction.reshape(rank, rank)) + + @staticmethod + def readable(attribute: Attribute) -> bool: + """Whether a header carries a full geometry — total, read-only, no I/O.""" + return all(key in attribute for key in _GEOMETRY_KEYS) + + @property + def rank(self) -> int: + return len(self.size_zyx) + + @property + def index_to_world(self) -> AffineMap: + """Continuous index ``(x, y, z)`` to world: ``p = O + D S i`` — ITK's own association.""" + return AffineMap(self.direction_xyz @ np.diag(self.spacing_xyz), np.asarray(self.origin_xyz, dtype=np.float64)) + + @property + def world_to_index(self) -> AffineMap: + return self.index_to_world.inverted() + + def _index_box(self, region_zyx: tuple[slice, ...] | None) -> tuple[np.ndarray, np.ndarray]: + """The region's outer faces as continuous indices ``(x, y, z)``: ``start - 0.5 .. stop - 0.5``. + + Outer faces and not voxel centres, deliberately: a sample is inside a grid while its + continuous index lies in ``[-0.5, n - 0.5)`` (the sampler rule of + ``Resample._resample_offset_region``), so a bound built on centres is short by the half + voxel that rule reaches. + """ + if region_zyx is None: + region_zyx = tuple(slice(0, extent) for extent in self.size_zyx) + low = np.array([float(part.start) - 0.5 for part in reversed(region_zyx)]) + high = np.array([float(part.stop) - 0.5 for part in reversed(region_zyx)]) + return low, high + + def world_box(self, region_zyx: tuple[slice, ...] | None = None) -> WorldBox: + """The axis-aligned world hull of a region's outer faces (the whole grid when ``None``).""" + low, high = self._index_box(region_zyx) + return WorldBox(low, high).image_under(self.index_to_world) + + def continuous_box(self, box: WorldBox) -> tuple[np.ndarray, np.ndarray]: + """A world box as a continuous-index box ``(low_xyz, high_xyz)`` on this grid.""" + image = box.image_under(self.world_to_index) + return image.low_xyz, image.high_xyz + + def index_window(self, box: WorldBox, margin: int) -> tuple[slice, ...]: + """The clamped array-order window a world box needs, grown by ``margin`` whole voxels. + + ``floor``/``ceil`` on the continuous-index box, plus the margin the interpolation taps + reach; clamped to a non-empty window, exactly as ``Resample._offset_window`` clamps — + a region entirely off the grid is a real place for a regrid (every sample takes the + fill) and a zero-width read is not something every backend serves. + """ + low, high = self.continuous_box(box) + window: list[slice] = [] + for axis in range(self.rank - 1, -1, -1): + start = int(np.floor(low[axis])) - margin + stop = int(np.ceil(high[axis])) + 1 + margin + extent = int(self.size_zyx[self.rank - 1 - axis]) + start = min(max(start, 0), extent - 1) + window.append(slice(start, min(max(stop, start + 1), extent))) + return tuple(window) + + def resampled( + self, + spacing_xyz: np.ndarray | None = None, + size_zyx: tuple[int, ...] | None = None, + align: str = "extent", + ) -> Grid: + """The same anatomy at another sampling density: give a spacing, or give a count. + + A component left at zero keeps that axis as it is. Whichever is given, the other follows, + and ``align`` decides where the new grid SITS — the one real choice here, worth a quarter + of a voxel of anatomy, and made silently by every library that offers only one of them: + + - ``extent`` — the outer faces coincide, so both grids cover exactly the same box and a + target index reads ``scale * (i + 0.5) - 0.5`` of the source. What ``F.interpolate`` does, + and what KonfAI has always done. + - ``origin`` — voxel zero's CENTRE stays put, so a target index reads ``scale * i`` and the + far edge moves by whatever the count rounded away. What resampling onto a grid that + shares an origin does. + + Under ``extent`` the spacing is derived from the counts and not from the request: a count is + a whole number, so the density that actually covers the box is ``n_src / n_dst`` times the + source's, and recording the requested one instead is a header that describes a grid nobody + sampled (up to a millimetre of drift across a volume, measured). + """ + if (spacing_xyz is None) == (size_zyx is None): + raise TransformError("A resampled grid is defined by a spacing or by a count, and by exactly one of them.") + if align not in ("extent", "origin"): + raise TransformError( + f"'{align}' is not a way to place a resampled grid.", + "Use align: extent to keep the field of view (the outer faces coincide) or align:" + " origin to keep voxel zero's centre where it is.", + ) + requested_xyz: np.ndarray | None = None + if size_zyx is not None: + counts = [ + int(want) if int(want) > 0 else int(have) for want, have in zip(size_zyx, self.size_zyx, strict=True) + ] + else: + wanted_xyz = np.asarray(spacing_xyz, dtype=np.float64) + requested_xyz = np.where(wanted_xyz > 0.0, wanted_xyz, self.spacing_xyz) + counts = [ + _voxel_count(have, source, wanted) + for have, source, wanted in zip(self.size_zyx, self.spacing_xyz[::-1], requested_xyz[::-1], strict=True) + ] + size = tuple(max(1, count) for count in counts) + # The density that actually covers the source's box with `size` voxels. + covering_xyz = self.spacing_xyz * np.array(self.size_zyx[::-1], dtype=np.float64) / np.array(size[::-1]) + if align == "origin": + return Grid( + size, self.origin_xyz, covering_xyz if requested_xyz is None else requested_xyz, self.direction_xyz + ) + return Grid( + size, + self.origin_xyz + 0.5 * (self.direction_xyz @ (covering_xyz - self.spacing_xyz)), + covering_xyz, + self.direction_xyz, + ) + + def sub_grid(self, region_zyx: tuple[slice, ...]) -> Grid: + """The grid of a region: same spacing and direction, the origin of its first voxel. + + The load-bearing line of every streamed regrid: a region left at the volume's origin + replays the volume's first slab wherever it lands, and the output still looks like an + image. The origin is ``index_to_world`` of the region's start — the same association of + the same product ITK uses in ``TransformIndexToPhysicalPoint``. + """ + start_xyz = np.array([float(part.start) for part in reversed(region_zyx)]) + return Grid( + tuple(int(part.stop - part.start) for part in region_zyx), + self.index_to_world.apply(start_xyz), + self.spacing_xyz, + self.direction_xyz, + ) + + +@dataclass(frozen=True) +class TransformBound: + """What a stored transform is guaranteed to do: an exact affine part and a bounded residual. + + ``T(p)`` lies in ``affine(p) ± residual_xyz`` for every ``p``, per world component. For a + linear transform the residual is zero and the statement is exact; for a BSpline it is the + sup-norm of the coefficients (non-negative basis functions summing to one make every + displacement a convex combination of them); for a dense field it is the recorded or declared + per-component bound. The affine part is read structurally off the transform, never probed: + a probe measures a local gradient and extrapolates it, which under-bounds (measured). + """ + + affine: AffineMap + residual_xyz: np.ndarray + + @staticmethod + def exact(affine: AffineMap) -> TransformBound: + return TransformBound(affine, np.zeros(affine.rank)) + + @staticmethod + def shift(residual_xyz: np.ndarray) -> TransformBound: + return TransformBound(AffineMap.identity(int(residual_xyz.size)), np.asarray(residual_xyz, dtype=np.float64)) + + def after(self, inner: TransformBound) -> TransformBound: + """The bound of ``self(inner(p))`` — interval arithmetic through the outer affine.""" + return TransformBound( + inner.affine.then(self.affine), + np.abs(self.affine.matrix) @ inner.residual_xyz + self.residual_xyz, + ) + + def map_box(self, box: WorldBox) -> WorldBox: + """Where the image of ``box`` is guaranteed to lie.""" + return box.image_under(self.affine).grown(self.residual_xyz) + + +@dataclass(frozen=True) +class AffineStage: + """One affine step of a decoded transform — exact, in world coordinates.""" + + map: AffineMap + + def bound(self) -> TransformBound: + return TransformBound.exact(self.map) + + +#: The B-spline orders KonfAI evaluates: the linear hat a dense field is read through, and the cubic +#: ITK writes a BSplineTransform with. ITK will happily write orders 0 and 2, which decode as +#: readily as any other and have no kernel here -- so the refusal belongs where the value is built, +#: not where it is finally sampled, which is mid-run and per region. +SUPPORTED_SPLINE_ORDERS = (1, 3) + + +@dataclass(frozen=True) +class DisplacementStage: + """One displacement step: ``p + d(p)``, with ``d`` interpolated off a value grid. + + One shape for the two non-linear things a stored transform can be. A BSpline is order-3 + coefficients on a coarse control grid; a dense field is order-1 samples on its own grid. Both + kernels are non-negative and sum to one, so the displacement anywhere is a convex combination + of ``values`` and ``sup |values|`` per component bounds it at every point — the bound that + replaces walking a region's boundary, which under-bounds a wiggle narrower than the region + and costs more than the resample it serves (both measured). + + ``values`` is ``(rank, Z, Y, X)`` float64, components in physical ``(x, y, z)`` order, world + units — ITK applies no direction matrix to them (verified in ``itkBSplineTransform.hxx``). + Outside the grid's reach the displacement is zero: ITK returns the identity there. + """ + + grid: Grid + values: np.ndarray + order: int + + def __post_init__(self) -> None: + if self.order not in SUPPORTED_SPLINE_ORDERS: + raise TransformError( + f"A displacement of B-spline order {self.order} is not one KonfAI evaluates" + f" (orders {', '.join(str(order) for order in SUPPORTED_SPLINE_ORDERS)}).", + "Write the transform as a displacement field, or as a cubic BSpline, which is what" + " every registration that produces one writes by default.", + ) + + @property + def bound_xyz(self) -> np.ndarray: + return np.abs(self.values.reshape(self.values.shape[0], -1)).max(axis=1) + + def bound(self) -> TransformBound: + return TransformBound.shift(self.bound_xyz) + + +#: A decoded stored transform: stages in APPLICATION order (first applied first). SimpleITK's +#: ``CompositeTransform`` applies its list in reverse — the decoder normalizes that here, once. +SpatialStages = tuple["AffineStage | DisplacementStage", ...] + + +def bound_of(stages: SpatialStages, rank: int) -> TransformBound: + """The bound of the whole decoded map, folded in application order.""" + folded = TransformBound.exact(AffineMap.identity(rank)) + for stage in stages: + folded = stage.bound().after(folded) + return folded diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 9ebf9a81..1808db96 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -37,7 +37,6 @@ LocalityKind, PatchLocality, RegionContext, - Resample, Save, Transform, split_expand, @@ -102,6 +101,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: ... def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -198,10 +198,14 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, ) -> list[slice]: + # A draw is bound to (case index, copy), not to the case's NAME: the name a region stage + # needs to find its own per-case map means nothing to an augmentation. + del name return self.augmentation.stream_region_source(self.index, self.a, target_slices, source_spatial_shape) def stream_region( @@ -228,7 +232,6 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) LocalityKind.HALO, LocalityKind.ORIENTATION, LocalityKind.CROP, - LocalityKind.RESCALE, LocalityKind.REGRID, ) @@ -254,25 +257,22 @@ def __call__(self, target: tuple[slice, ...]) -> list[slice]: @dataclass(frozen=True) class _RemapPull: - """An index-remap stage's pull map, bound to the case state the stages before it left.""" - - remap: Callable[[tuple[slice, ...], list[int], Attribute], list[slice]] - shape: list[int] - attribute: Attribute - - def __call__(self, target: tuple[slice, ...]) -> list[slice]: - return self.remap(target, list(self.shape), Attribute(self.attribute)) - + """An index-remap stage's pull map, bound to the case and the state the stages before it left. -@dataclass(frozen=True) -class _ScalePull: - """A rescale stage's pull map: the source window of the scale mapping plus its interpolation halo.""" + The case NAME is bound here because a stage instance is shared by every case of a manager + (``DatasetManager`` hands the same transforms list to each), while a map read from a stored + transform or a reference header is per case. A pull that could not say which case it was for + would build one case's window from another case's map -- and a window that is short does not + raise, it returns the fill. + """ - scales: list[float] + remap: Callable[[str, tuple[slice, ...], list[int], Attribute], list[slice]] shape: list[int] + attribute: Attribute + name: str = "" def __call__(self, target: tuple[slice, ...]) -> list[slice]: - return Resample.source_window(target, self.scales, self.shape) + return self.remap(self.name, target, list(self.shape), Attribute(self.attribute)) @dataclass(frozen=True) @@ -1736,13 +1736,13 @@ def _plan_stream_region( in hand at the moment each check fails, so the reason is built here -- dropping it would leave every fallback silent. ``evolved`` is the case state the plan leaves, which a :class:`Save` sweep writes as its cache header. The chain streams when every - stage is pointwise, a region kind (``HALO``/``ORIENTATION``/``CROP``/``RESCALE`` — any + stage is pointwise, a region kind (``HALO``/``ORIENTATION``/``CROP``/``REGRID`` — any number, each pulling through the one before it), or a ``GLOBAL_STAT`` with a pre-populated statistic. The plan walks the chain once with one evolving case state, so each stage declares against — and remaps from — the geometry the stages before it left, and folds the spatial shapes stage by stage; a fold that does not land on ``landing_shape`` (the copy's own grid by default) refuses (the safety net for a stage whose shape map is not declared). Any - ``WHOLE_VOLUME`` declaration, an unreadable ``GLOBAL_STAT``, a ``RESCALE`` without a known + ``WHOLE_VOLUME`` declaration, an unreadable ``GLOBAL_STAT``, a ``REGRID`` without a known ``Spacing`` (or that is not a :class:`Resample`), or a halo too wide to be worth reading rejects streaming. ``seed_statistics=False`` accepts a missing statistic instead of reading it — for a chain fed by a cache that is not materialized yet, whose re-resolution seeds it @@ -1793,17 +1793,6 @@ def refuse(reason: str) -> tuple[bool, tuple[_ReadStagePlan, ...], Attribute, st f"{label} declares a halo of {loc.halo} that is too wide for this grid to be worth" " reading (over half the patch extent per axis)." ) - if loc.kind is LocalityKind.RESCALE and (not isinstance(stage, Resample) or "Spacing" not in evolved): - # A resample is patch-native only when the source geometry is known: the scale is read - # from the evolving 'Spacing' (a free geometry stat, no read_data_statistics). - return refuse( - f"{label} declares RESCALE but " - + ( - "does not inherit from Resample." - if not isinstance(stage, Resample) - else "the source carries no 'Spacing' to scale from." - ) - ) plan = self._plan_read_stage(stage, loc, shape, evolved) plans.append(plan) shape = list(plan.out_shape) @@ -1826,14 +1815,8 @@ def _plan_read_stage( return _ReadStagePlan( loc.kind, tuple(shape), tuple(shape), _HaloPull(_halo_radii(loc.halo, len(shape)), list(shape)) ) - if loc.kind is LocalityKind.RESCALE: - resample = cast(Resample, stage) - out = [int(e) for e in resample.transform_shape(self.group_src, self.name, list(shape), Attribute(evolved))] - scales = [shape[k] / out[k] for k in range(len(shape))] - resample.write_stream_cache_attribute(evolved, list(shape)) - return _ReadStagePlan(loc.kind, tuple(shape), tuple(out), _ScalePull(scales, list(shape))) - # ORIENTATION / CROP: the stage's own remap, evaluated on the state the stages before it left. - pull = _RemapPull(stage.stream_region_source, list(shape), Attribute(evolved)) + # ORIENTATION / CROP / REGRID: the stage's own remap, on the state the stages before it left. + pull = _RemapPull(stage.stream_region_source, list(shape), Attribute(evolved), self.name) out = self._stage_out_shape(stage, shape, Attribute(evolved)) stage.write_stream_cache_attribute(evolved, list(shape)) return _ReadStagePlan(loc.kind, tuple(shape), tuple(out), pull) @@ -2744,7 +2727,7 @@ def _get_streamed_data( def _finalize_stream_patch(self, tensor: torch.Tensor, index: int, a: int, is_input: bool) -> torch.Tensor: """Pad/format a target-extent streamed patch to ``patch_size`` like the whole-volume path. - The region and RESCALE streamed paths produce a patch at the raw target-slice extent, but the + The region streamed paths produce a patch at the raw target-slice extent, but the overlap tiling can leave the last patch narrower than ``patch_size`` (integer-floor stride). The whole-volume ``Patch.get_data`` pads that border patch up to ``patch_size`` via ``apply_read_plan``; running the streamed patch through the SAME read plan makes border patches @@ -2808,7 +2791,7 @@ def _replay_streamed_region( volume) and crops it back after the stage — the stage's own edge padding reproduces the whole-volume border once the clamp reaches the true border, so seams agree. ORIENTATION applies the index remap to what its region read; a CROP's remap IS its action, so the stage is - not re-applied; RESCALE interpolates the sub-region to its target extent. Only the composed + not re-applied; REGRID interpolates the sub-region to its target extent. Only the composed region is requested; whether that avoids decoding the whole volume depends on the storage format -- compressed MetaImage and NRRD decode the full volume per read (see ``_supports_region_read``). @@ -2849,15 +2832,7 @@ def _replay_streamed_region( # a throwaway scope, and write the case-level answer once from the FULL shape below # (write_stream_cache_attribute). scoped = Attribute(cache_attribute) - if plan.kind is LocalityKind.RESCALE: - tensor = cast(Resample, stage).resample_region( - tensor, - tuple(target), - [s.start for s in source], - [plan.in_shape[k] / plan.out_shape[k] for k in range(len(plan.in_shape))], - list(plan.in_shape), - ) - elif plan.kind is not LocalityKind.CROP: + if plan.kind is not LocalityKind.CROP: # A HALO stage is handed the ENLARGED region it asked for, and told so: what it # returns is cropped back to the target just below. tensor = stage.stream_region( diff --git a/konfai/data/sampling.py b/konfai/data/sampling.py new file mode 100644 index 00000000..c3f10d43 --- /dev/null +++ b/konfai/data/sampling.py @@ -0,0 +1,454 @@ +# 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 + +"""Where each target voxel reads from, and the gather that reads it — in torch, on the volume's device. + +Two halves, and only the first one varies. The COORDINATE PRODUCER turns a decoded transform into +one source index per target voxel; the GATHER is ITK's sampler and is written once. A stage that +resamples onto another grid, one that warps through a field and one that does both differ only in +what they hand the producer. + +Nothing here calls SimpleITK. A BSpline is evaluated from its coefficient grid and a dense field +from its samples, with the same tensor-product kernel arithmetic ITK uses (verified to 8e-14 +against ``TransformPoint``), so a chain that resamples through a stored transform stays on the GPU +from the read to the write. +""" + +from __future__ import annotations + +import itertools + +import numpy as np +import torch + +from konfai.data.geometry import ( + SUPPORTED_SPLINE_ORDERS, + AffineMap, + AffineStage, + DisplacementStage, + Grid, + SpatialStages, + TransformBound, +) + +#: Coordinates are accumulated in float64 and only the gather runs in the payload's dtype. A world +#: coordinate is an origin of hundreds of millimetres plus an index of hundreds of voxels, and in +#: float32 that sum keeps about four digits past the voxel -- enough to move a sample across a +#: voxel boundary on a large grid, which is a visibly different value on anything that is not +#: smooth. Measured on a 64x512x512 slab: float32 coordinates disagree with SimpleITK by 0.35 on +#: data in [0, 1], float64 by 3e-05 (which is the gather's own float32 rounding). +_COORDINATE_DTYPE = torch.float64 + + +def _kernel_weights(offset: torch.Tensor, order: int) -> torch.Tensor: + """The 1-D B-spline weight at distance ``offset`` — ITK's own kernels. + + Order 1 is the linear hat; order 3 is ``itkBSplineKernelFunction``'s cubic. Both are + non-negative and sum to one over their support, which is what makes ``sup |values|`` a bound + on the displacement at every point rather than at the samples. + """ + distance = offset.abs() + if order == 1: + return torch.clamp(1.0 - distance, min=0.0) + if order == 3: + near = (4.0 - 6.0 * distance**2 + 3.0 * distance**3) / 6.0 + far = (2.0 - distance) ** 3 / 6.0 + return torch.where(distance < 1.0, near, torch.where(distance < 2.0, far, torch.zeros_like(distance))) + # Unreachable: DisplacementStage refuses any other order where it is built. + raise ValueError(f"No B-spline kernel of order {order}: KonfAI evaluates {SUPPORTED_SPLINE_ORDERS}.") + + +def _apply(points_xyz: torch.Tensor, affine: AffineMap, device: torch.device) -> torch.Tensor: + """``translation + Σ_j column_j · p_j``, accumulated in ``j`` order — ITK's own association. + + Not ``points @ matrix.T + offset``. That is the same sum in a different order, which BLAS is + free to reassociate, and the two answers differ in the last bit; a continuous index landing on + an exact half then rounds to the other voxel. The one matmul saved is not worth a label map + that disagrees with ``sitk.Resample`` in whole columns. + """ + matrix = torch.tensor(affine.matrix, dtype=_COORDINATE_DTYPE, device=device) + out = torch.tensor(affine.translation, dtype=_COORDINATE_DTYPE, device=device).expand(points_xyz.shape).clone() + for j in range(affine.rank): + out = out + points_xyz[..., j, None] * matrix[:, j] + return out + + +def _to_index(world_xyz: torch.Tensor, grid: Grid, device: torch.device) -> torch.Tensor: + """World to continuous index, in ITK's arithmetic: subtract the origin FIRST, then accumulate. + + ``TransformPhysicalPointToContinuousIndex`` sums ``M[i][j] * (p_j - O_j)`` from zero. Folding the + origin into a translation instead — which is what a generic inverted affine map is — reassociates + a difference of large world coordinates against a small one and moves the last bit. + """ + matrix = torch.tensor(grid.world_to_index.matrix, dtype=_COORDINATE_DTYPE, device=device) + origin = torch.tensor(grid.origin_xyz, dtype=_COORDINATE_DTYPE, device=device) + shifted = world_xyz - origin + out = torch.zeros_like(world_xyz) + for j in range(grid.rank): + out = out + shifted[..., j, None] * matrix[:, j] + return out + + +def _displacement_at(stage: DisplacementStage, world_xyz: torch.Tensor, device: torch.device) -> torch.Tensor: + """The stage's displacement at each world point — zero where its grid does not reach. + + ITK returns the identity outside a BSpline's valid region and outside a field's domain, so a + point past the values simply moves by nothing. That is also what makes a region of a field a + legal stand-in for the whole of it: the part it does not cover was going to be identity here + anyway — which is exactly why the region handed in must COVER the target region, and why a + short one degrades in silence rather than raising. + """ + rank = stage.grid.rank + index = _to_index(world_xyz, stage.grid, device) # continuous index on the value grid, (x, y, z) + values = torch.tensor(stage.values, dtype=_COORDINATE_DTYPE, device=device) + extent_xyz = [int(stage.grid.size_zyx[rank - 1 - axis]) for axis in range(rank)] + + base = torch.floor(index) - (stage.order - 1) // 2 + taps = stage.order + 1 + # The two domains ITK actually implements, and they differ. A BSpline is the identity unless its + # whole support lies in the coefficient grid (``InsideValidRegion``), so its edge falls off a + # control point early. A dense field is an ordinary image: it interpolates anywhere in + # ``[-0.5, n - 0.5)`` with the taps clamped, which reaches half a voxel past the outermost + # samples. Using the spline's rule for a field blanks that rim -- 10% of the voxels of a small + # volume, all of them at the border, all of them silently un-warped. + inside = torch.ones(index.shape[:-1], dtype=torch.bool, device=device) + for axis in range(rank): + if stage.order == 1: + inside &= (index[..., axis] >= -0.5) & (index[..., axis] < extent_xyz[axis] - 0.5) + else: + inside &= (base[..., axis] >= 0) & (base[..., axis] + taps - 1 < extent_xyz[axis]) + + flat_values = values.reshape(rank, -1) + out = torch.zeros(index.shape, dtype=_COORDINATE_DTYPE, device=device) + for corner in itertools.product(range(taps), repeat=rank): + weight = torch.ones(index.shape[:-1], dtype=_COORDINATE_DTYPE, device=device) + positions_xyz = [] + for axis in range(rank): + position = base[..., axis] + corner[axis] + weight = weight * _kernel_weights(index[..., axis] - position, stage.order) + positions_xyz.append(position.to(torch.long).clamp(0, extent_xyz[axis] - 1)) + # ``values`` is (component, Z, Y, X) and row-major over the spatial axes, so the flat index + # runs the ARRAY axes outermost-first with x fastest -- the mirror of the physical order the + # coordinates arrive in. Running it the other way samples the field transposed, which warps + # the anatomy somewhere plausible and wrong. + flat_index = torch.zeros(index.shape[:-1], dtype=torch.long, device=device) + for array_axis in range(rank): + flat_index = flat_index * int(stage.grid.size_zyx[array_axis]) + positions_xyz[rank - 1 - array_axis] + gathered = torch.stack([flat_values[c].index_select(0, flat_index.reshape(-1)) for c in range(rank)], dim=-1) + out = out + gathered.reshape(out.shape) * weight.unsqueeze(-1) + return out * inside.unsqueeze(-1) + + +def source_index( + target_grid: Grid, + source_grid: Grid, + stages: SpatialStages, + device: torch.device, +) -> torch.Tensor: + """One source continuous index per voxel of ``target_grid``, shaped ``(*size_zyx, rank)`` in ``(x, y, z)``. + + ``target_grid`` is the REGION's grid — its own origin, not the volume's. + + THE WORLD POINT IS MATERIALISED, never folded away. Target index to world and world to source + index compose into one affine that is algebraically identical and one matmul cheaper, and it is + the wrong arithmetic: ITK takes the two steps separately, so the two associations disagree in the + last bit — and a continuous index landing EXACTLY on ``k + 0.5`` (which is what a target grid + commensurate with its source produces, in whole columns at a time) then rounds to a different + voxel. Measured against ``sitk.Resample``: 150 of 1050 voxels of a label map, every one of them + a legal-looking label. Consecutive affine STAGES are still folded — that is between two world + points, where nothing rounds to an index. + """ + rank = target_grid.rank + axes = [ + torch.arange(int(extent), dtype=_COORDINATE_DTYPE, device=device) for extent in reversed(target_grid.size_zyx) + ] + # meshgrid in array order (Z, Y, X) with the physical components last, so the tensor indexes + # like the volume it will sample. + grids = torch.meshgrid(*reversed(axes), indexing="ij") + world = _apply(torch.stack(list(reversed(grids)), dim=-1), target_grid.index_to_world, device) + + pending = AffineMap.identity(rank) + for stage in stages: + if isinstance(stage, AffineStage): + pending = pending.then(stage.map) + continue + if not pending.is_identity: + world = _apply(world, pending, device) + pending = AffineMap.identity(rank) + world = world + _displacement_at(stage, world, device) + if not pending.is_identity: + world = _apply(world, pending, device) + return _to_index(world, source_grid, device) + + +def _is_diagonal(matrix: np.ndarray) -> bool: + """Whether every off-diagonal entry is EXACTLY zero — no tolerance, deliberately. + + A tolerance would admit maps whose separable form is only nearly the general one, and the two + would then disagree in the last bit at exactly the coordinates that round to a different voxel. + Exact is also not restrictive: an axis-aligned or axis-flipped grid gives exact zeros, and the + inverse of such a matrix keeps them. + """ + return not np.any(matrix - np.diag(np.diag(matrix))) + + +def separable_source_index( + target_grid: Grid, + source_grid: Grid, + stages: SpatialStages, + device: torch.device, +) -> list[torch.Tensor] | None: + """One source index per target ROW of each array axis, or ``None`` when the map does not factorise. + + THE SAME ARITHMETIC, with the terms that are exactly zero left out. Each component of + :func:`source_index` is ``translation_k + Σ_j p_j · M[k, j]``; with ``M`` diagonal every ``j ≠ k`` + contributes ``p_j · 0.0``, and adding an exact zero changes no float. So this is bit-identical + where it applies rather than merely close -- which is what lets one case take it and another the + general path without the two ever having to be reconciled. + + What it saves is not arithmetic but MEMORY TRAFFIC: the general path materialises a coordinate + per voxel and gathers eight corners through a flat index over the whole volume. + """ + if stages: + return None + forward, backward = target_grid.index_to_world, source_grid.world_to_index + if not (_is_diagonal(forward.matrix) and _is_diagonal(backward.matrix)): + return None + rank = target_grid.rank + axes: list[torch.Tensor] = [] + for array_axis in range(rank): + axis = rank - 1 - array_axis # the physical component this array axis runs along + index = torch.arange(int(target_grid.size_zyx[array_axis]), dtype=_COORDINATE_DTYPE, device=device) + world = float(forward.translation[axis]) + index * float(forward.matrix[axis, axis]) + axes.append((world - float(source_grid.origin_xyz[axis])) * float(backward.matrix[axis, axis])) + return axes + + +def blend_order(target_grid: Grid, source_grid: Grid) -> list[int]: + """The array axes in the order to blend them: most source voxels per target voxel first. + + Each axis is reduced to its output extent before the next one reads it, so blending the axis + that shrinks most FIRST makes every later pass smaller. On an isotropic change of spacing that + is worth nothing; on a thick-slice CT brought to isotropic -- 64x512x512 at 3x0.7x0.7 mm, where + z triples while y and x shrink -- it is 44 M intermediate elements against 110 M. + + Keyed on the two grids' SPACINGS, never on the extents in hand: a streamed region and the whole + volume have to blend in the same order or they stop being bit-identical, and their extents + differ by construction while their spacings do not. + """ + rank = target_grid.rank + scale = [ + abs(float(target_grid.spacing_xyz[rank - 1 - axis] / source_grid.spacing_xyz[rank - 1 - axis])) + for axis in range(rank) + ] + return sorted(range(rank), key=lambda array_axis: -scale[array_axis]) + + +def gather_separable( + source: torch.Tensor, + axes: list[torch.Tensor], + source_starts_zyx: list[int], + source_shape_zyx: list[int], + mode: str, + fill: float, + blend: list[int] | None = None, +) -> torch.Tensor: + """:func:`gather`'s rules over a map that factorises — one ``index_select`` per axis, no volume. + + Same interval, same tap clamp, same round-half-up, same fill. It sums the tensor product in a + different ORDER than the general gather -- axis by axis rather than corner by corner -- so the + two agree to float rounding rather than bit for bit. That costs nothing: a map either factorises + or it does not, so no case is ever served by both. The equality that has to be exact is a + streamed region against the whole volume, and it is: the per-axis coordinates are global, so a + region takes a sub-range of the very numbers the whole volume takes. + """ + from konfai.data.transform import nearest_index, sampling_dtype, window_index + + rank = len(axes) + window_zyx = [int(extent) for extent in source.shape[1:]] + extent_zyx = [int(axis.numel()) for axis in axes] + device = source.device + + inside_axes = [(axis >= -0.5) & (axis < source_shape_zyx[array_axis] - 0.5) for array_axis, axis in enumerate(axes)] + out_shape = [int(source.shape[0]), *extent_zyx] + if not all(bool(mask.any()) for mask in inside_axes): + return torch.full(out_shape, fill, device=device, dtype=torch.float32).type(source.dtype) + + def local(index: torch.Tensor, array_axis: int) -> torch.Tensor: + return window_index(index, source_shape_zyx[array_axis], source_starts_zyx[array_axis], window_zyx[array_axis]) + + def broadcast(values: torch.Tensor, array_axis: int) -> torch.Tensor: + shape = [1] * (rank + 1) + shape[array_axis + 1] = -1 + return values.reshape(shape) + + def is_identity(array_axis: int) -> bool: + """Whether this axis reads itself: the same voxels, in order, unblended.""" + axis = axes[array_axis] + if int(axis.numel()) != window_zyx[array_axis] or source_starts_zyx[array_axis] != 0: + return False + return bool(torch.equal(axis, torch.arange(int(axis.numel()), dtype=axis.dtype, device=device))) + + out = source if mode == "nearest" else source.type(sampling_dtype(source)) + for array_axis in range(rank) if blend is None else blend: + # An axis the map leaves alone is read by leaving it alone: two gathers and a blend that + # would reproduce the input exactly, over the largest tensor in flight, for nothing. + if is_identity(array_axis): + continue + axis = axes[array_axis] + if mode == "nearest": + out = out.index_select(array_axis + 1, local(nearest_index(axis), array_axis)) + continue + # One axis at a time, not eight corners at once. The tensor product is the same sum, but + # blending axis by axis reduces each extent before the next axis reads it -- six gathers over + # shrinking tensors instead of twenty-four over the largest one. + base = torch.floor(axis) + share = broadcast((axis - base).to(out.dtype), array_axis) + index = base.to(torch.long) + low = out.index_select(array_axis + 1, local(index, array_axis)) + high = out.index_select(array_axis + 1, local(index + 1, array_axis)) + # lerp fuses the three passes `low * (1 - w) + high * w` into one. Exact at w = 0, which is + # the only endpoint reachable: w is `x - floor(x)`, so it never reaches 1. + out = torch.lerp(low, high, share) + if mode == "nearest" and not out.is_floating_point(): + out = out.type(torch.float32) + + if all(bool(mask.all()) for mask in inside_axes): + # Nothing of this region falls outside the source, which is the ordinary case for a resample + # within a volume -- so neither the mask nor the pass that applies it is built at all. + return out.type(source.dtype) + mask = inside_axes[0] + for array_axis in range(1, rank): + mask = mask.unsqueeze(-1) & inside_axes[array_axis] + return out.masked_fill(~mask.unsqueeze(0), fill).type(source.dtype) + + +def gather( + source: torch.Tensor, + coordinates_xyz: torch.Tensor, + source_starts_zyx: list[int], + source_shape_zyx: list[int], + mode: str, + fill: float, +) -> torch.Tensor: + """ITK's sampler at an arbitrary coordinate per voxel, over the region actually read. + + The rule is ``sitk.Resample``'s, and it is the same one the separable samplers in + ``konfai.data.transform`` obey: a sample is inside while its continuous source index lies in + ``[-0.5, n - 0.5)``; inside, the interpolation taps clamp to the buffer, so the half-voxel rim + past the outermost voxel centres reproduces the border value instead of falling off; outside, + the sample is ``fill``. + + ``source`` covers ``source_starts_zyx`` onward of a volume of ``source_shape_zyx``; the + coordinates are GLOBAL indices of that volume. + + THE TWO MODES TAKE DIFFERENT ROUTES, and the difference is what each can afford. Nearest is one + gather on the exact index, because the pick is discontinuous and the last bit decides it. Linear + is eight gathers over a flat index, which ``grid_sample`` does as one fused kernel -- at the cost + of normalising by the extent it is handed, so a streamed region and the whole volume agree to + ~1e-5 rather than exactly. + """ + rank = coordinates_xyz.shape[-1] + window_zyx = [int(extent) for extent in source.shape[1:]] + extent_zyx = list(coordinates_xyz.shape[:-1]) + device = source.device + + inside = torch.ones(extent_zyx, dtype=torch.bool, device=device) + for axis in range(rank): + array_axis = rank - 1 - axis + coordinate = coordinates_xyz[..., axis] + inside &= (coordinate >= -0.5) & (coordinate < source_shape_zyx[array_axis] - 0.5) + + out_shape = [int(source.shape[0]), *extent_zyx] + if not bool(inside.any()): + return torch.full(out_shape, fill, device=device, dtype=torch.float32).type(source.dtype) + + from konfai.data.transform import nearest_index, sampling_dtype, window_index + + work = source.type(sampling_dtype(source)) + if mode == "nearest": + # One gather, on exact index arithmetic: a nearest pick is discontinuous, so the last bit of + # a coordinate decides which voxel it lands on. Measured over 300 random grid pairs, a + # single-precision coordinate moves 2.6% of the picks -- in a label map, 2.6% of the voxels + # wearing a different label. There are no eight corners to fuse here, so nothing is bought + # by handing it to a kernel that normalises coordinates. + flat = torch.zeros(extent_zyx, dtype=torch.long, device=device) + for array_axis in range(rank): + index = nearest_index(coordinates_xyz[..., rank - 1 - array_axis]) + local = window_index( + index, source_shape_zyx[array_axis], source_starts_zyx[array_axis], window_zyx[array_axis] + ) + flat = flat * window_zyx[array_axis] + local + picked = work.reshape(int(work.shape[0]), -1).index_select(1, flat.reshape(-1)).reshape(out_shape) + return picked.masked_fill(~inside.unsqueeze(0), fill).type(source.dtype) + + # A BLEND goes through grid_sample, which is one fused kernel where the eight corners are eight + # gathers over a flat index. `align_corners=False` IS ITK's domain -- a sample is inside while + # its continuous index lies in [-0.5, n - 0.5) -- and `padding_mode="border"` IS ITK's tap clamp; + # what grid_sample has no notion of is the FILL, so the mask computed above still applies it. + # + # It costs the bit-identity between a streamed region and the whole volume, and that is the whole + # of what it costs: grid_sample takes NORMALISED coordinates, so it divides by the extent of the + # tensor handed to it, and a region is handed a window. The two agree to ~1e-5 instead of + # exactly. Deliberate, and asked for -- 4x on a warp. + local_axes = [] + for array_axis in range(rank): + axis = rank - 1 - array_axis + extent = window_zyx[array_axis] + shifted = coordinates_xyz[..., axis] - float(source_starts_zyx[array_axis]) + local_axes.append((2.0 * shifted + 1.0) / extent - 1.0) + # grid_sample orders the last dimension (x, y, z) -- the mirror of the array axes, which is the + # order the coordinates already arrive in. + sampling = torch.stack([local_axes[rank - 1 - axis] for axis in range(rank)], dim=-1).to(work.dtype) + out = torch.nn.functional.grid_sample( + work.unsqueeze(0), sampling.unsqueeze(0), mode="bilinear", padding_mode="border", align_corners=False + ).squeeze(0) + return out.masked_fill(~inside.unsqueeze(0), fill).type(source.dtype) + + +def source_window( + target_grid: Grid, + source_grid: Grid, + bound: TransformBound, + margin: int = 1, +) -> tuple[slice, ...]: + """The source window a target region pulls, from the bound alone — no voxel sampled. + + Closed form, which is what lets a cost model price a decomposition without doing the bounding + work for it: the target region's world box, mapped through the bound (an exact affine hull + grown by the residual), read back as a clamped index window on the source. + """ + return source_grid.index_window(bound.map_box(target_grid.world_box()), margin) + + +def read_amplification( + target_grid: Grid, + source_grid: Grid, + bound: TransformBound, + regions: list[tuple[slice, ...]], + margin: int = 1, +) -> float: + """How many times the source's voxels this decomposition reads, in total. + + The honest name for what streaming costs: every region's source window is read whole, the + windows overlap, and the finer the decomposition the more they overlap. Monotone in fineness, + so it is a property of the decomposition and not of the transform alone. + """ + total = 0 + for region in regions: + window = source_window(target_grid.sub_grid(region), source_grid, bound, margin) + total += int(np.prod([part.stop - part.start for part in window], dtype=np.int64)) + return float(total) / float(np.prod(source_grid.size_zyx, dtype=np.int64)) diff --git a/konfai/data/transform.py b/konfai/data/transform.py index bfa2ed52..addc46f7 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -16,7 +16,6 @@ """Tensor and image transforms used in KonfAI preprocessing and postprocessing.""" -import itertools import os import tempfile from abc import ABC, abstractmethod @@ -36,6 +35,23 @@ import torch.nn.functional as F from konfai import cuda_visible_devices +from konfai.data.geometry import ( + _GEOMETRY_KEYS, + AffineMap, + AffineStage, + DisplacementStage, + Grid, + SpatialStages, + TransformBound, +) +from konfai.data.sampling import ( + blend_order, + gather, + gather_separable, + separable_source_index, + source_index, + source_window, +) from konfai.utils.config import _escape_key_component, apply_config from konfai.utils.dataset import Attribute, Dataset, DataStream, data_to_image, image_to_data from konfai.utils.errors import TransformError @@ -60,11 +76,9 @@ class LocalityKind(Enum): box, so it is no bijection and the stored volume's statistics are not its output's. - ``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. + - ``REGRID`` -- resample onto another grid: a change of sampling density, of placement, or + both, possibly through a map. The target is a grid in its own right, so part of it may read + from outside the source altogether and the source region is no mere scaling of the target's. 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 @@ -79,7 +93,6 @@ class LocalityKind(Enum): ORIENTATION = "orientation" CROP = "crop" GLOBAL_STAT = "global_stat" - RESCALE = "rescale" REGRID = "regrid" SLAB = "slab" WHOLE_VOLUME = "whole_volume" @@ -91,7 +104,7 @@ def preserves_statistics(self) -> bool: Only a reorientation does: a flip or a permute is a bijection on the voxels, so the multiset of values -- and therefore Min/Max/Mean/Std over it -- is exactly the input's. Every other kind may map values (``POINTWISE``, ``GLOBAL_STAT``), mix neighbours (``HALO``) or interpolate - (``RESCALE``). This is what decides whether the statistics of the STORED volume are still those + (``REGRID``). This is what decides whether the statistics of the STORED volume are still those of a later transform's own input (see ``DatasetManager._plan_stream_region``). """ return self is LocalityKind.ORIENTATION @@ -191,6 +204,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -198,9 +212,9 @@ def stream_region_source( """Map a target-patch's spatial slices to the source spatial region to read (region kinds). Overridden by the kinds whose source region is an index remap of the target's -- ``ORIENTATION`` - maps it and reorients what it reads, ``CROP`` maps it and is done. ``HALO`` and ``RESCALE`` are - handled generically by the dispatcher, so the base raises for any transform that declares a - region kind without providing the remap. + maps it and reorients what it reads, ``CROP`` maps it and is done, ``REGRID`` maps it through + its own geometry. ``HALO`` is handled generically by the dispatcher, so the base raises for + any other transform that declares a region kind without providing the remap. ``cache_attribute`` is the case's SOURCE metadata, under the same rules as :meth:`patch_locality`: a remap the image decides (a reorientation whose mirrored axes are the @@ -323,7 +337,7 @@ def inverse_patch_locality(self, cache_attribute: Attribute) -> PatchLocality: The default derives from the forward contract where the derivation is safe for any subclass: a per-voxel value map inverts to a per-voxel value map, and an index remap inverts to an index remap. Every other kind falls to ``WHOLE_VOLUME`` — an inverse that is streamable anyway - (``Padding``'s crop, ``Resample``'s rescale) declares itself. + (``Padding``'s crop, ``Resample``'s change of grid) declares itself. """ forward = self.patch_locality(cache_attribute) if forward.kind in (LocalityKind.POINTWISE, LocalityKind.ORIENTATION): @@ -340,8 +354,36 @@ def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) """ return shape + def inverse_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + """State the attribute transition ``inverse`` makes, instead of performing it. + + The write mirror of :meth:`write_stream_cache_attribute`, and it exists because the streamed- + write dispatcher plans a pipe by walking a ONE-VOXEL probe through it: a stage whose inverse + restores a whole volume cannot be run on that probe just to learn what it pops. The base is a + no-op -- an inverse that pops nothing has nothing to state, and one whose transition is cheap + to perform is simply run. + """ + + def stream_region_inverse( + self, + name: str, + tensor: torch.Tensor, + context: RegionContext, + cache_attribute: Attribute, + ) -> torch.Tensor: + """Apply ``inverse`` to a region, told WHERE that region sits — the mirror of + :meth:`Transform.stream_region`. + + ``context.target`` is the region of the inverse's OUTPUT being produced and ``context.source`` + the region of its input on hand. The default delegates to :meth:`inverse`, so an involutive + index remap (whose pulled block already IS the answer's input) keeps working untouched. + """ + del context + return self.inverse(name, tensor, cache_attribute) + def stream_region_target( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -851,6 +893,7 @@ def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) def stream_region_target( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -940,1144 +983,875 @@ def window_index(index: torch.Tensor, n_in: int, region_start: int, window: int) return torch.clamp(torch.clamp(index, 0, n_in - 1) - region_start, 0, window - 1) -class Resample(TransformInverse, ABC): - def __init__(self, inverse: bool) -> None: - super().__init__(inverse) +# --------------------------------------------------------------------------------------------- +# One resample. Two questions: which grid to write on, and what map to write it through. +# --------------------------------------------------------------------------------------------- - def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: - # The source region is derived from the scale mapping (read from cache_attribute['Spacing'] - # by the dispatcher); a small interpolation halo is added by resample_source_region. - return PatchLocality(LocalityKind.RESCALE) - def _resample(self, tensor: torch.Tensor, size: list[int]) -> torch.Tensor: - if tensor.dtype == torch.uint8: - mode = "nearest" - elif len(tensor.shape) < 4: - mode = "bilinear" - else: - mode = "trilinear" +class _TargetGrid(ABC): + """Which grid a resample writes on — the ``to`` half of the question.""" - work = tensor.type(sampling_dtype(tensor)) - # Return on the input's device (interpolate preserves it): a CPU input stays on the CPU, a - # GPU-resident output volume stays on the GPU so the whole finalize runs where the volume is. - return F.interpolate(work.unsqueeze(0), size=tuple(size), mode=mode).squeeze(0).type(tensor.dtype) + #: The geometry keys this target cannot be built without. An extent change needs none; a + #: density change needs the Spacing; adopting another grid needs a real physical space. + needs: frozenset[str] = frozenset() @abstractmethod - def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - pass + def of(self, source: Grid, name: str) -> Grid: + """The grid a case stored on ``source`` is written on.""" + + def set_datasets(self, datasets: list[Dataset]) -> None: # noqa: B027 - only a reference has one + """The run's roots, for a target that has an image of its own to look up.""" @abstractmethod - def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: - pass + def describe(self) -> str: + """The target named as a refusal or a plan line names it.""" - def _inverse_geometry(self, cache_attribute: Attribute) -> list[int]: - """Pop the Size/Spacing stack the forward pushed and return the size the inverse restores.""" - cache_attribute.pop_np_array("Size") - size_1 = cache_attribute.pop_np_array("Size") - if "Spacing" in cache_attribute: - cache_attribute.pop_np_array("Spacing") - return [int(size) for size in size_1] - def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - return self._resample(tensor, self._inverse_geometry(cache_attribute)) +class _OwnGrid(_TargetGrid): + """No change of grid: the map moves what the voxels hold, not where they are.""" - def inverse_patch_locality(self, cache_attribute: Attribute) -> PatchLocality: - # The inverse rescales back to the size the forward stacked: patch-native (RESCALE) whenever - # that stack is on the finalize-time attribute, judged on a copy (a declaration never pops). - try: - self._inverse_geometry(Attribute(cache_attribute)) - except NameError: - return PatchLocality(LocalityKind.WHOLE_VOLUME) - return PatchLocality(LocalityKind.RESCALE) + def of(self, source: Grid, name: str) -> Grid: + del name + return source - def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) -> list[int]: - try: - return self._inverse_geometry(Attribute(cache_attribute)) - except NameError: - return shape + def describe(self) -> str: + return "the case's own grid" - def stream_region_target( - self, - target_slices: tuple[slice, ...], - source_spatial_shape: list[int], - cache_attribute: Attribute, - ) -> list[slice]: - # The inverse rescales the accumulator (n_in) back to the stored size: a written region pulls - # through the same coordinate formula as the forward read, with the roles swapped. - n_in = [int(s) for s in source_spatial_shape] - n_out = self.inverse_transform_shape(list(n_in), cache_attribute) - scales = [n_in[k] / n_out[k] for k in range(len(n_in))] - return Resample.source_window(target_slices, scales, n_in) - - # Every patch derives its source coordinates from the same global scale (n_in / n_out, from the - # truncated integer sizes F.interpolate itself uses), which is what makes the streamed patches - # agree with the whole-volume call and with each other across a seam. - #: What this stage interpolates with, or ``None`` to read it off the dtype. A subclass taking an - #: ``interpolation`` argument assigns it here, and every sampler asks the one method below -- - #: a declaration honoured on one path and not another is worse than none, because the page that - #: tells a user to set it is then right about half the chains. - interpolation: str | None = None - - def _stream_mode(self, tensor: torch.Tensor) -> str: - """``nearest``, or the rank's linear name -- what a sampler asks before it blends anything. - A dtype cannot settle this on its own: a CT is int16 and so is nothing else about it. The - heuristic therefore claims ``uint8`` and nothing more, and a stage exposing ``interpolation`` - answers for everything it cannot know. Getting it wrong is silent -- two blended labels give - a third that was in no input, in a volume that is still a label map. - """ - declared = self.interpolation or ("nearest" if tensor.dtype == torch.uint8 else "linear") - if declared == "nearest": - return "nearest" - return "bilinear" if len(tensor.shape) < 4 else "trilinear" +class _DerivedGrid(_TargetGrid): + """The case's own grid at another density — a spacing, or a count, and where it sits.""" - def resample_source_region( - self, - target_slices: tuple[slice, ...], - source_spatial_shape: list[int], - cache_attribute: Attribute, - halo: int = 1, - ) -> tuple[list[slice], list[int], list[float], list[int], list[int]]: - """Map a TARGET-grid patch to the minimal SOURCE region to read. + def __init__(self, spacing: list[float] | None, shape: list[int] | None, align: str) -> None: + self.spacing = None if spacing is None else np.asarray([max(0.0, float(value)) for value in spacing]) + self.shape = None if shape is None else tuple(max(0, int(value)) for value in shape) + self.align = align + # A density is meaningless without the density it starts from; a count is not. + self.needs = frozenset({"Spacing"}) if spacing is not None else frozenset() - Returns ``(source_slices, region_starts, scales, n_in, n_out)`` — all in - array axis order (Z, Y, X). The ``halo`` is a pure safety margin (the - formula's ``+2`` already captures the i1 neighbour); nearest needs none. - """ - n_in = [int(s) for s in source_spatial_shape] - n_out = [int(s) for s in self.transform_shape("", "", list(n_in), cache_attribute)] - scales = [n_in[k] / n_out[k] for k in range(len(n_in))] - source_slices = Resample.source_window(target_slices, scales, n_in, halo) - return source_slices, [s.start for s in source_slices], scales, n_in, n_out + def of(self, source: Grid, name: str) -> Grid: + where = f"case '{name}'" if name else "the case" + if self.spacing is not None: + if self.spacing.size != source.rank: + raise TransformError( + f"'Resample' was given a spacing of {self.spacing.size} value(s) and {where} has" + f" {source.rank} spatial axis/axes." + ) + return source.resampled(spacing_xyz=self.spacing, align=self.align) + shape = cast("tuple[int, ...]", self.shape) + if len(shape) != source.rank: + raise TransformError( + f"'Resample' was given a shape of {len(shape)} value(s) and {where} has" + f" {source.rank} spatial axis/axes." + ) + return source.resampled(size_zyx=shape, align=self.align) - @staticmethod - def source_window( - target_slices: tuple[slice, ...] | list[slice], - 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. - - Covers BOTH samplers, because the same window serves either mode: the linear taps around the - half-pixel source (``scale * (o + 0.5) - 0.5``, plus the ``+2``/``halo`` margin for the i1 - neighbour) AND the voxel nearest picks (``floor(o * scale)`` -- F.interpolate's own nearest - 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)) - smax = int(np.floor(scales[k] * ((sl.stop - 1) + 0.5) - 0.5)) - near_lo = int(np.floor(sl.start * scales[k])) - near_hi = int(np.floor((sl.stop - 1) * scales[k])) - start = min(smin - halo, near_lo) - stop = max(smax + 2 + halo, near_hi + 1) - source_slices.append(slice(max(0, start), min(n_in[k], stop))) - return source_slices + def describe(self) -> str: + if self.spacing is not None: + return f"a spacing of {[float(value) for value in self.spacing]}" + return f"a shape of {list(cast('tuple[int, ...]', self.shape))}" - @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 +class _ReferenceGrid(_TargetGrid): + """The grid of a STORED image: extent, spacing, origin and direction, read from its header. - def resample_region( - self, - sub_tensor: torch.Tensor, - target_slices: tuple[slice, ...], - 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. + The target that makes a cohort foldable. A spacing lines up densities and a shape lines up + extents, but both leave each case where it was; this adopts one grid whole, which is what gives + ``Reduce``'s ``grid: strict`` something true to compare. That is the atlas-template build. - ``sub_tensor`` is ``[C, (z, y, x)]`` covering ``source_slices``; - ``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]``. + THE REFERENCE IS AN IMAGE, NOT A LIST OF NUMBERS. A grid is fifteen numbers in two axis orders + at once, 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 it: 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. + """ - ``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) - if mode == "nearest": - indices = [] - for k in range(ndim): - # Take the axis's index map from F.interpolate itself, so streamed nearest picks the - # same source voxel as the whole-volume call for every size ratio. - src = torch.arange(n_in[k], device=dev, dtype=torch.float32).reshape(1, 1, -1) - n_out_k = round(n_in[k] / scales[k]) - index = F.interpolate(src, size=n_out_k, mode="nearest").long().flatten() - indices.append(index[target_slices[k].start : target_slices[k].stop] - region_starts[k]) - # One gather over broadcast index views instead of one volume copy per axis (nearest is a - # pure coordinate gather, so composing the axes changes no value). - return sub_tensor[(slice(None), *torch.meshgrid(*indices, indexing="ij"))] - - work = sub_tensor.type(sampling_dtype(sub_tensor)) - taps: list[tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]] = [] - for k in range(ndim): - o = torch.arange(target_slices[k].start, target_slices[k].stop, device=dev, dtype=work.dtype) - src = torch.clamp(scales[k] * (o + 0.5) - 0.5, min=0.0) - i0 = torch.floor(src).long() - i1 = torch.clamp(i0 + 1, max=n_in[k] - 1) - lam = src - i0.to(work.dtype) - taps.append(((i0 - region_starts[k], 1 - lam), (i1 - region_starts[k], lam))) - out_shape = [work.shape[0]] + [sl.stop - sl.start for sl in target_slices] - out = torch.zeros(out_shape, device=dev, dtype=work.dtype) - for combo in itertools.product(*taps): - gathered = work - weight = torch.ones([1] * (ndim + 1), device=dev, dtype=work.dtype) - for k, (idx, lam) in enumerate(combo): - gathered = gathered.index_select(k + 1, idx) - shape = [1] * (ndim + 1) - shape[k + 1] = -1 - weight = weight * lam.reshape(shape) - 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], - fill: float | None = None, - ) -> 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. - - ``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:]] - # 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, outside, device=device, dtype=torch.float32).type(sub_tensor.dtype) - - def local(index: torch.Tensor, k: int) -> torch.Tensor: - return window_index(index, n_in[k], region_starts[k], window[k]) - - if self._stream_mode(sub_tensor) == "nearest": - picks = [local(nearest_index(axis), 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: - work = sub_tensor.type(sampling_dtype(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), outside).type(sub_tensor.dtype) + needs = frozenset(_GEOMETRY_KEYS) - @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. + def __init__(self, entry: str, group: str | None, dataset: str | None) -> None: + self.entry = str(entry).strip() + self.group = group + # A root of its own, or the run's: left out, the grid to adopt is one member of the very + # cohort being brought together. + self.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.dataset = Dataset(Path(filename), file_format) + self.roots: list[Dataset] = [] + self._grid: Grid | None = None - Called once per case on the persistent attribute so ``inverse()`` at - prediction time pops exactly what the non-streamed path pushed. Uses the - FULL source shape, never the halo'd sub-region. - """ + def set_datasets(self, datasets: list[Dataset]) -> None: + self.roots = list(datasets) + def _roots(self) -> list[Dataset]: + return [self.dataset] if self.dataset is not None else list(self.roots) -class ResampleToResolution(Resample): - def __init__(self, spacing: list[float] = [1.0, 1.0, 1.0], inverse: bool = True) -> None: - super().__init__(inverse) - self.spacing = torch.tensor([0 if s < 0 else s for s in spacing]) + 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"'Resample' cannot tell which group of '{dataset.filename}' holds reference" + f" '{self.entry}': it has {len(groups)} ({', '.join(sorted(groups)) or 'none'}).", + "Name it: Resample: {reference: " + self.entry + ", reference_group: }.", + ) - def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: - if "Spacing" not in cache_attribute: + def grid(self) -> Grid: + """The reference's grid, read from its header once. + + 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( - "Missing 'Spacing' in cache attributes, the data is likely not a valid image.", - "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.", + f"'Resample' has no dataset to look reference '{self.entry}' up in.", + "Give the stage a root of its own -- Resample: {reference: " + + self.entry + + ", reference_dataset: ./Reference:omezarr} -- or run it in a workflow, which hands" + " its dataset_filenames to every stage.", ) - if len(shape) != len(self.spacing): - raise TransformError(f"Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}") - image_spacing = cache_attribute.get_tensor("Spacing") - resize_factor = torch.tensor( - [s / i_s if s > 0 else 1.0 for s, i_s in zip(self.spacing, image_spacing, strict=False)] + 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) + self._grid = Grid.of([int(extent) for extent in shape[1:]], attribute, f"reference '{self.entry}'") + return self._grid + raise TransformError( + f"'Resample' cannot find reference '{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.", ) - return [int(x) for x in (torch.tensor(shape) * 1 / resize_factor.flip(0))] - def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - image_spacing = cache_attribute.get_tensor("Spacing") - spacing = self.spacing - resize_factor = torch.tensor( - [ - s / i_s if s > 0 else 1.0 - for s, i_s in zip(self.spacing, cache_attribute.get_tensor("Spacing"), strict=False) - ] - ) - cache_attribute["Spacing"] = torch.tensor( - [float(s) if s > 0 else float(i_s) for s, i_s in zip(spacing, image_spacing, strict=False)] - ) - cache_attribute["Size"] = np.asarray([int(x) for x in torch.tensor(tensor.shape[1:])]) - size = [int(x) for x in (torch.tensor(tensor.shape[1:]) * 1 / resize_factor.flip(0))] - cache_attribute["Size"] = np.asarray(size) - return self._resample(tensor, size) + def of(self, source: Grid, name: str) -> Grid: + grid = self.grid() + if grid.rank != source.rank: + where = f"case '{name}'" if name else "the case" + raise TransformError( + f"'Resample' cannot resample {where}, which has {source.rank} spatial axis/axes," + f" onto reference '{self.entry}', which has {grid.rank}." + ) + return grid - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: - image_spacing = cache_attribute.get_tensor("Spacing") - spacing = self.spacing - resize_factor = torch.tensor( - [s / i_s if s > 0 else 1.0 for s, i_s in zip(self.spacing, image_spacing, strict=False)] - ) - cache_attribute["Spacing"] = torch.tensor( - [float(s) if s > 0 else float(i_s) for s, i_s in zip(spacing, image_spacing, strict=False)] - ) - cache_attribute["Size"] = np.asarray([int(x) for x in source_spatial_shape]) - size = [int(x) for x in (torch.tensor([int(s) for s in source_spatial_shape]) * 1 / resize_factor.flip(0))] - cache_attribute["Size"] = np.asarray(size) + def describe(self) -> str: + return f"reference '{self.entry}'" -class ResampleToShape(Resample): - def __init__(self, shape: list[float] = [100, 256, 256], inverse: bool = True) -> None: - super().__init__(inverse) - self.shape = torch.tensor([0 if s < 0 else s for s in shape]) +class Resample(TransformInverse): + """Resample a case: onto another grid, through a stored map, or both — in one interpolation. - def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: - if "Spacing" not in cache_attribute: - raise TransformError( - "Missing 'Spacing' in cache attributes, the data is likely not a valid image.", - "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.", - ) - if len(shape) != len(self.shape): - raise TransformError(f"Shape and target dimensions do not match: shape={shape}, target_shape={self.shape}") - new_shape = self.shape.clone() - for i, s in enumerate(self.shape): - if s == 0: - new_shape[i] = shape[i] - return new_shape + Every resample in KonfAI is these two questions, and this is the only stage that answers them. - def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - shape = self.shape.clone() - image_shape = torch.tensor([int(x) for x in torch.tensor(tensor.shape[1:])]) - for i, s in enumerate(self.shape): - if s == 0: - shape[i] = image_shape[i] - if "Spacing" in cache_attribute: - cache_attribute["Spacing"] = torch.flip( - image_shape / shape * torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), - dims=[0], - ) - cache_attribute["Size"] = image_shape - cache_attribute["Size"] = shape - return self._resample(tensor, shape) + **Which grid to write on** — at most one of: - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: - shape = self.shape.clone() - image_shape = torch.tensor([int(s) for s in source_spatial_shape]) - for i, s in enumerate(self.shape): - if s == 0: - shape[i] = image_shape[i] - if "Spacing" in cache_attribute: - cache_attribute["Spacing"] = torch.flip( - image_shape / shape * torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), - dims=[0], - ) - cache_attribute["Size"] = image_shape - cache_attribute["Size"] = shape + - nothing (the default): the case's own grid. The map moves the anatomy; the voxels stay put. + - ``spacing``: the same field of view at another density. A component left at ``0`` keeps its axis. + - ``shape``: the same field of view at a given count. A component left at ``0`` keeps its axis. + - ``reference``: the grid of a stored image, adopted whole — extent, spacing, origin, direction. + **What map to write it through** — any of, composed in this order: -@dataclass(frozen=True) -class _ReferenceMap: - """Everything a case needs to be read onto a reference grid, computed once from the headers. + - ``field``: a displacement field, read in world units at each TARGET voxel. Its own grid, its + own spacing: a field solved at 120 um moves a volume stored at 30 um without being upsampled. + - ``transforms``: transforms stored beside the cases — rigid, affine, BSpline, dense field, or a + composite of them — mapping GROUP to whether to invert it. The LAST declared is applied first, + which is SimpleITK's own composite order. - 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. - """ + Left out, the map is the identity and this is a change of grid and nothing else. - 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 + ONE INTERPOLATION, ALWAYS. A grid change and a warp asked for together are composed into a single + coordinate per target voxel and the source is read once, at the displaced point. Doing it as two + stages resamples twice, and a volume interpolated twice has lost detail the second pass invented + no more of -- which is the whole reason an atlas's appearance is rebuilt from native volumes. + IT STREAMS, and what a region reads is known before a voxel is touched. A rigid or affine map is + an exact affine, so the source box of a target region is that region's box mapped through it. A + BSpline and a dense field are values on a grid read through a non-negative kernel that sums to + one, so the sup-norm of those values bounds the displacement at EVERY point -- a theorem, not a + sample of the boundary. A field on disk is bounded by ``max_displacement`` instead, which is then + CHECKED against every region actually read. -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. - - ``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, 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; - - 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. A field with - no bound at all declares ``WHOLE_VOLUME`` and says so, as ``Warp`` does. + ``align`` decides where a ``spacing`` or a ``shape`` grid SITS, and it is the one silent choice + in the family -- a quarter of a voxel of anatomy, made differently by every library that offers + only one of them. ``extent`` keeps the field of view (the outer faces coincide); ``origin`` keeps + voxel zero's centre where it is. A ``reference`` states its own placement and ignores this. + + WHAT IT REFUSES, rather than resample from a window it cannot size or in a space it does not have: + + - a case whose header carries no ``Origin``/``Spacing``/``Direction`` when the answer needs + physical space (a reference, a stored transform, a field). A plain ``spacing``/``shape`` + resample does not: with no geometry a world coordinate IS an index, and the ratio is the map; + - a transform type that decomposes into no bounded map, naming the type; + - ``invert: true`` on anything but a rigid or affine map: inverting a spline or a field is a + dense solve over the whole grid, and a field solved per region is not the restriction of the + field solved once. Store the inverse, or invert it where it is written; + - a field with no ``max_displacement`` to size its region from; + - a case that does not meet the target grid anywhere -- the output would be ``fill`` from edge to + edge, and an all-background member is a plausible, wrong contribution to a median. + + Every refusal but the last declares ``WHOLE_VOLUME`` with its reason and the run proceeds on the + whole-volume path, so a chain never breaks over one: it only stops being bounded, and says so in + the plan. A case reaching only PART of the target grid is legal and common -- the rest takes + ``fill`` -- and the plan prints how much of the grid it covers. """ def __init__( self, - entry: str, - group: str | None = None, - dataset: str | None = None, + spacing: list[float] | None = None, + shape: list[int] | None = None, + reference: str | None = None, + reference_group: str | None = None, + reference_dataset: str | None = None, + transforms: dict[str, bool] | None = None, field: str | None = None, field_group: str | None = None, max_displacement: float | str = 0.0, - fill: float = 0.0, + align: str = "extent", interpolation: str | None = None, + fill: float = 0.0, inverse: bool = True, ) -> None: super().__init__(inverse) if interpolation is not None and interpolation not in ("linear", "nearest"): raise TransformError( - f"'ResampleToReference' has an unknown interpolation '{interpolation}'.", + f"'Resample' has an unknown interpolation '{interpolation}'.", "Use 'linear' for an image or 'nearest' for a label map. Left unset, uint8 is taken" " for a label map and everything else is interpolated.", ) self.interpolation = interpolation - if not entry or not str(entry).strip(): + self.fill_value = float(fill) + self._target = self._target_from(spacing, shape, reference, reference_group, reference_dataset, align) + if transforms is not None and not transforms: raise TransformError( - "'ResampleToReference' needs an 'entry': the stored image whose grid to adopt.", - "Name it, e.g. ResampleToReference: {entry: 822174, group: Volume}.", + "'Resample' was given an empty 'transforms'.", + "Name a group and say whether to invert it -- transforms: {reg: false} -- or drop the" + " argument: without it the map is the identity and this is a change of grid alone.", ) - 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) - # 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. + self.transforms = transforms 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.", + f"'Resample' 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: it sizes the" + " region a field is read from and means nothing without one.", ) self.displacement: _DisplacementSource | None = ( - _DisplacementSource( - "ResampleToReference", field, field_group, max_displacement, group_keyword="field_group" - ) + _DisplacementSource("Resample", field, field_group, max_displacement, group_keyword="field_group") 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, _ReferenceMap] = {} + #: Per case: the grid its own header describes. Recorded where that header is in hand -- + #: transform_shape, called for every case as the manager is built. A region read hands back + #: the REGION's Origin, so a grid rebuilt from what a streamed region arrives with would + #: 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. + self._grids: dict[str, Grid] = {} + #: Per case: the geometry keys its header did not carry (see :meth:`Grid.from_header`). + self._assumed: dict[str, frozenset[str]] = {} + self._stored: dict[str, SpatialStages] = {} + self._refusal: str | None = None + self._probed = False + + @staticmethod + def _target_from( + spacing: list[float] | None, + shape: list[int] | None, + reference: str | None, + reference_group: str | None, + reference_dataset: str | None, + align: str, + ) -> _TargetGrid: + named = [name for name, value in (("spacing", spacing), ("shape", shape), ("reference", reference)) if value] + if len(named) > 1: + raise TransformError( + f"'Resample' was given {' and '.join(named)}, which are three ways to say the same thing.", + "A resample writes on one grid: give its density (spacing), its extent (shape) or the" + " image whose grid to adopt (reference) -- and only one of them.", + ) + if align not in ("extent", "origin"): + raise TransformError( + f"'Resample' has an unknown align '{align}'.", + "Use align: extent to keep the field of view (the outer faces coincide, which is what" + " KonfAI has always done) or align: origin to keep voxel zero's centre where it is.", + ) + if reference: + return _ReferenceGrid(reference, reference_group, reference_dataset) + if spacing is not None or shape is not None: + return _DerivedGrid(spacing, shape, align) + if reference_group is not None or reference_dataset is not None: + raise TransformError( + "'Resample' was told where to find a reference but not which one.", + "Name the entry whose grid to adopt: Resample: {reference: 822174, reference_group: Volume}.", + ) + return _OwnGrid() def set_datasets(self, datasets: list[Dataset]) -> None: super().set_datasets(datasets) + self._target.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 + # ------------------------------------------------------------------ the two grids - def _roots(self) -> list[Dataset]: - return [self.reference_dataset] if self.reference_dataset is not None else list(self.datasets) + @property + def _needs(self) -> frozenset[str]: + """The geometry keys this configuration cannot be answered without. - 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: }.", - ) + A stored map or a reference grid is applied in physical space and needs all three; a change + of density needs the density it starts from; a change of extent needs nothing at all. Being + exact about this is what lets one class serve a headerless array and a real volume. + """ + if self.transforms is not None or self.displacement is not None: + return frozenset(_GEOMETRY_KEYS) + return self._target.needs - 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. + def _record(self, name: str, shape: list[int], cache_attribute: Attribute) -> Grid: + """The case's own grid, remembered under its name, with what its header left unsaid.""" + where = f"case '{name}'" if name else "the case" + grid, missing = Grid.from_header(list(shape), cache_attribute, where) + self._assumed[name] = missing + if name: + self._grids[name] = grid + return grid - 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: + def _source_grid(self, name: str) -> Grid: + grid = self._grids.get(name) + if grid is None: 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.", + f"'Resample' 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 grid of every case as" + " its manager is built, and a region is only ever streamed afterwards.", ) - 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.", - ) + return grid - @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: + def _target_of(self, name: str) -> tuple[Grid, Grid]: + """``(source, target)`` — needs only what BUILDING the target grid needs. + + Split from :meth:`_grids_of` because the two questions have different answers: the output + SHAPE of a warp on the case's own grid is the case's own shape, knowable with no geometry at + all, while SAMPLING it is not. Refusing the shape too would take down the plan of a chain + whose honest answer is to fall back to the whole volume and say so. + """ + source = self._source_grid(name) + absent = self._assumed.get(name, frozenset()) + lacking = [key for key in _GEOMETRY_KEYS if key in absent and key in self._target.needs] + if lacking: 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" + f"'Resample' cannot place {self._target.describe()} for case '{name}': its header" + f" carries no {', '.join(lacking)}.", + "A density is meaningless without the density it starts from, and another grid" + " cannot be adopted without a physical space to adopt 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): + return source, self._target.of(source, name) + + def _grids_of(self, name: str) -> tuple[Grid, Grid]: + source = self._source_grid(name) + absent = self._assumed.get(name, frozenset()) + lacking = [key for key in _GEOMETRY_KEYS if key in absent and key in self._needs] + if lacking: 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.", + f"'Resample' needs the geometry of case '{name}' to resample it onto" + f" {self._target.describe()}, and its header carries no {', '.join(lacking)}.", + "Resampling onto another grid, or through a stored map, 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).", ) - return origin, spacing, direction.reshape(rank, rank) + return source, self._target.of(source, name) # ------------------------------------------------------------------ the map - @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. - - 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" - 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) - 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] = recorded - return recorded + def _stored_stages(self, name: str) -> SpatialStages: + """This case's stored transforms, decoded and composed, in application order.""" + if name in self._stored: + return self._stored[name] + from konfai.utils.ITK import decode_transform_stages, invert_stages - @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.", - ) + _require_simpleitk() + rank = self._source_grid(name).rank + stages: list[AffineStage | DisplacementStage] = [] + # Reversed: a CompositeTransform applies its members last-first, and this stage has always + # built one from `transforms` in declaration order. Decoding normalizes each member to + # application order, so the declared list is reversed here to mean the same thing it did. + for group in reversed(list(cast("dict[str, bool]", self.transforms))): + invert = self.transforms[group] if self.transforms else False + stored = None + for dataset in self.datasets: + if dataset.is_dataset_exist(group, name): + stored = dataset.read_transform(group, name) + break + if stored is None: + raise TransformError( + f"'Resample' found no transform for case '{name}' in group '{group}'.", + "Every case needs an entry in every group named under 'transforms:'. Check the" + " group name, or drop the cases that have no transform with 'subset'.", + ) + decoded = decode_transform_stages(stored) + if invert: + inverted = invert_stages(decoded, rank) + if inverted is None: + raise TransformError( + f"'Resample' cannot invert group '{group}' for case '{name}': it is not a rigid or affine map.", + "Inverting a spline or a displacement field is a dense solve over the whole" + " grid, and a field solved per region is not the restriction of the field" + f" solved once. Store the inverse field instead, or set '{group}: false' and" + " invert it where it is written.", + ) + decoded = inverted + stages.extend(decoded) + self._stored[name] = tuple(stages) + return self._stored[name] - 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 _field_stage(self, name: str, region: Grid) -> DisplacementStage: + """The declared field over ``region``, read on its own grid and no wider. - 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. + The field is evaluated at the TARGET's world points, so the window it needs is that region's + own world box -- no halo, whatever the displacement is. What the halo sizes is the SOURCE + read, which is a different question answered by the bound. """ - if self.displacement is None or not name: - return None, None, None - shape, attribute = self.displacement.infos(name) + source = cast("_DisplacementSource", self.displacement) + shape, attribute = source.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 - 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 - 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 ( - f"case '{name}' covers {covered * 100:.1f}% of reference '{self.entry}';" - f" the rest of what it writes is fill ({self.fill_value:g})" - ) + 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)) + source.check_bound(values, name) + return DisplacementStage(grid.sub_grid(window), values.numpy(), order=1) + + def _stages(self, name: str, region: Grid) -> SpatialStages: + """The whole map over one target region, in application order.""" + stages: list[AffineStage | DisplacementStage] = [] + if self.displacement is not None: + stages.append(self._field_stage(name, region)) + if self.transforms is not None: + stages.extend(self._stored_stages(name)) + return tuple(stages) + + def _bound(self, name: str) -> TransformBound: + """What the map is guaranteed to do — from declarations and coefficients, no voxel read.""" + rank = self._source_grid(name).rank + folded = TransformBound.exact(AffineMap.identity(rank)) + if self.displacement is not None: + declared = self.displacement.component_bound() + if declared is None: + raise TransformError(self.displacement.undeclared_reason()) + folded = TransformBound.shift(np.asarray(declared[:rank], dtype=np.float64)).after(folded) + if self.transforms is not None: + for stage in self._stored_stages(name): + folded = stage.bound().after(folded) + return folded # ------------------------------------------------------------------ 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).target + del group_src + self._record(name, [int(extent) for extent in shape], cache_attribute) + _source, target = self._target_of(name) + if name: + self._refuse_if_disjoint(name) + return [int(extent) for extent in target.size_zyx] 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: + # The geometry is judged on the attribute in hand -- the case's own header, as the base + # contract has it -- and not on what the cohort has been seen to carry: one case of a group + # may lack an Origin while the rest have one, and a declaration made per case is the honest + # one. A config-time probe hands over an empty header, which reads as a case with none. + lacking = [key for key in _GEOMETRY_KEYS if key in self._needs and key not in cache_attribute] + if lacking: 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" + f"resampling onto {self._target.describe()} happens in physical space and this" + f" case carries no {', '.join(lacking)}. Use a source whose geometry is readable" + " (mha, nii, h5, or an OME-Zarr written by KonfAI)" ), ) + if not self._probed: + self._probed = True + self._refusal = self._probe_cohort() + if self._refusal is not None: + return PatchLocality(LocalityKind.WHOLE_VOLUME, reason=self._refusal) 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. + def _probe_cohort(self) -> str | None: + """Whether every case this stage will see is boundable, or the sentence saying which is not. - 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. + The COHORT's answer, not one case's: a locality is declared once for the stage while the + cases are many, so a group whose entries are not uniformly decodable must fall back for all + of them rather than for the ones that happen to be planned first. Exceptions are swallowed + into a reason -- this runs inside the plan, where a raise would take the run down instead of + costing it the whole-volume path. The GEOMETRY is not judged here; that is per case, and + :meth:`patch_locality` reads it off the header it is handed. """ - 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)))) + if self.transforms is not None and sitk is None: + return ( + "SimpleITK is not installed, and a stored transform is applied in physical space by" + " it. Install it (pip install konfai[itk]) to stream this stage" + ) + # The field's bound is the COHORT's, read from declarations and headers, so it is answered + # before any case has been seen -- and it is what a config-time probe is really asking. + if self.displacement is not None and self.displacement.component_bound() is None: + return self.displacement.undeclared_reason() + for name in self._grids: + try: + self._bound(name) + except TransformError as error: + # Both halves of the refusal: the first says what is wrong, the second what to + # change. A plan line carrying only the first tells the reader nothing to do. + return " ".join(str(part).strip() for part in error.args if part) + except Exception: # an unreadable transform is a whole-volume answer, not a crash + return ( + f"the map for case '{name}' could not be read, so what it does to a region is" + " unknown. Check the group names under 'transforms:'/'field:' and that every" + " case has an entry in each" + ) + return None def stream_region_source( - self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + self, name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute ) -> list[slice]: - shape = [int(extent) for extent in source_spatial_shape] - 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, - ) + del source_spatial_shape, cache_attribute + source, target = self._grids_of(name) + return list(source_window(target.sub_grid(tuple(target_slices)), source, self._bound(name))) 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(). + # The recorded grid, not one read off `cache_attribute`: what arrives here describes the + # REGION, down to an Origin of its own. See _record(). del cache_attribute - return self._sample_target_region( - name, - tensor, - self._recorded(name), - tuple(context.target), - [sl.start for sl in context.source], - [int(extent) for extent in context.source_shape], - ) + return self._sample(name, tensor, tuple(context.target), [part.start for part in context.source]) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: shape = [int(extent) for extent in tensor.shape[1:]] - recorded = self.grid_map(name, shape, cache_attribute) + if name not in self._grids: + self._record(name, shape, cache_attribute) + source, target = self._grids_of(name) # 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 - ) + whole = tuple(slice(0, extent) for extent in target.size_zyx) + result = self._sample(name, tensor, whole, [0] * source.rank) 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], + def _sample( + self, name: str, sub_tensor: torch.Tensor, target_slices: tuple[slice, ...], region_starts: list[int] ) -> torch.Tensor: - """One region of the target grid, read from the source in a SINGLE interpolation. + source, target = self._grids_of(name) + region = target.sub_grid(target_slices) + stages = self._stages(name, region) + shape, mode = list(source.size_zyx), self._mode(sub_tensor) + # A map that factorises is read one axis at a time, which is the same arithmetic without the + # terms that are zero and without a coordinate per voxel -- and it is most maps, because most + # volumes are stored axis-aligned. The general form is what a rotation or a displacement + # needs, and the two are bit-identical wherever both apply. + axes = separable_source_index(region, source, stages, sub_tensor.device) + if axes is not None: + order = blend_order(target, source) + return gather_separable(sub_tensor, axes, region_starts, shape, mode, self.fill_value, order) + coordinates = source_index(region, source, stages, sub_tensor.device) + return gather(sub_tensor, coordinates, region_starts, shape, mode, self.fill_value) + + def _mode(self, tensor: torch.Tensor) -> str: + """``nearest`` or ``linear`` — what a sampler asks before it blends anything. - 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. + A dtype cannot settle this on its own: a CT is int16 and so is nothing else about it. The + heuristic therefore claims ``uint8`` and nothing more, and ``interpolation`` answers for + everything it cannot know. Getting it wrong is silent -- two blended labels give a third + that was in no input, in a volume that is still a label map. """ - 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) + declared = self.interpolation or ("nearest" if tensor.dtype == torch.uint8 else "linear") + return "nearest" if declared == "nearest" else "linear" + + def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + """Push the target grid over the source's, so the case now IS the grid it was written on. + + Pushed and not replaced: the source geometry stays underneath for :meth:`inverse` to pop back + to, which is the whole of the stack this and ``_inverse_geometry`` share. + """ + shape = [int(extent) for extent in source_spatial_shape] + source, missing = Grid.from_header(shape, cache_attribute, "the case") + target = self._target.of(source, "") + written = { + "Spacing": target.spacing_xyz, + "Origin": target.origin_xyz, + "Direction": target.direction_xyz.ravel(), + } + for key in _GEOMETRY_KEYS: + # Only over a geometry that was there: a case stored without an Origin is resampled by + # ratio, and inventing one for it would be a header nobody measured. Key by key, and not + # all-or-nothing, because ``inverse`` pops exactly what is present -- so what is pushed + # and what is popped are the same condition, read at the two ends. + if key not in missing: + cache_attribute[key] = written[key] + cache_attribute["Size"] = np.asarray(shape) + cache_attribute["Size"] = np.asarray([int(extent) for extent in target.size_zyx]) + + # ------------------------------------------------------------------ the plan + + #: Below this, a case is worth a line in the plan: it reaches only part of the target 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 + + #: How many probes per axis the coverage estimate uses. Coverage is a volume ratio between two + #: boxes that a rotation makes a polytope, so it is counted rather than solved; capped because + #: it is a plan line, not a result. + _COVERAGE_PROBES = 24 + + def coverage(self, name: str) -> float: + """The fraction of the target grid that reads from inside the recorded case.""" + return self._coverage(*self._target_of(name)) + + @classmethod + def _coverage(cls, source: Grid, target: Grid) -> float: + """The fraction of ``target`` that reads from inside ``source``, from geometry alone. + + The GRID CHANGE only, never the map: a stored transform is what makes the two meet, and + counting the target as uncovered because a registration has not been applied yet would call + every warp disjoint. 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. + """ + axes = [ + np.linspace(0.0, float(extent) - 1.0, min(cls._COVERAGE_PROBES, int(extent))) + for extent in reversed(target.size_zyx) ] - 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. + lattice = np.stack([axis.ravel() for axis in np.meshgrid(*axes, indexing="ij")], axis=-1) + index = target.index_to_world.then(source.world_to_index).apply(lattice) + inside = np.ones(index.shape[0], dtype=bool) + for axis in range(source.rank): + inside &= (index[:, axis] >= -0.5) & (index[:, axis] < source.size_zyx[source.rank - 1 - axis] - 0.5) + return float(np.count_nonzero(inside)) / float(inside.size) + + def _refuse_if_disjoint(self, name: str) -> None: + """Refuse a case that does not meet the target 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. """ - 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, + if self._target_is_own or self.coverage(name) > 0.0: + return + where = f"case '{name}'" if name else "the case" + raise TransformError( + f"'Resample' would write {where} as nothing but 'fill': it does not overlap" + f" {self._target.describe()} anywhere, so no voxel of the target 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 target the cohort" + " actually surrounds, or drop this case with 'subset'.", ) - # (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.type(sampling_dtype(sub_tensor)) - flat_source = work.reshape(int(work.shape[0]), -1) - - def _tap(offsets: list[torch.Tensor]) -> torch.Tensor: - flat_index = torch.zeros(extent, dtype=torch.long, device=sub_tensor.device) - for axis, offset in enumerate(offsets): - index = window_index(offset, n_in[axis], region_starts[axis], window[axis]) - flat_index = flat_index * window[axis] + index - return flat_source.index_select(1, flat_index.reshape(-1)).reshape(out_shape) - - if self._stream_mode(sub_tensor) == "nearest": - picked = _tap([nearest_index(axis).expand(extent) for axis in coordinates]) - return picked.masked_fill(~inside.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) - out = torch.zeros(out_shape, device=sub_tensor.device, dtype=work.dtype) - for corner in itertools.product((0, 1), repeat=rank): - weight = torch.ones(extent, device=sub_tensor.device, dtype=work.dtype) - for axis, step in enumerate(corner): - weight = weight * (weights[axis] if step else 1 - weights[axis]).to(work.dtype) - out += _tap([bases[axis].to(torch.long) + step for axis, step in enumerate(corner)]) * weight - return out.masked_fill(~inside.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) + 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. - 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) + Not on the grid recorded for the case: the plan asks a stage about its own input, which the + stages before it decide, and a note answered from the stored header would describe a volume + that no longer exists by the time this stage sees it. Nothing is recorded here either, for + the mirror reason -- a question must not move the state a region read depends on. + """ + del group_dest + try: + source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") + if missing & self._target.needs: + return None + covered = self._coverage(source, self._target.of(source, name)) + except TransformError: + return None + if covered >= self._WORTH_SAYING: + return None + return ( + f"case '{name}' covers {covered * 100:.1f}% of {self._target.describe()};" + f" the rest of what it writes is fill ({self.fill_value:g})" + ) # ------------------------------------------------------------------ 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 + """Pop the geometry stack the forward pushed and return the size the inverse restores.""" + cache_attribute.pop_np_array("Size") + size = cache_attribute.pop_np_array("Size") + for key in _GEOMETRY_KEYS: + # Present iff the forward pushed it (see write_stream_cache_attribute): popping restores + # the case's own, and a key the case never had is one this never wrote. + if key in cache_attribute: + cache_attribute.pop_np_array(key) + return [int(extent) for extent in size] + + @staticmethod + def _grid_from(cache_attribute: Attribute, shape: list[int]) -> Grid: + if Grid.readable(cache_attribute): + return Grid.of(shape, cache_attribute, "the case") + return Grid.identity(shape) + + def _inverse_grids(self, cache_attribute: Attribute, shape: list[int]) -> tuple[Grid, Grid]: + """``(what the accumulator is on, what to write back onto)`` — both off the pushed stack. + + The forward stacked the source geometry under the target's, so the inverse needs no memory + of the case: it reads the grid it is holding, pops, and reads the grid it is restoring. A + copy is popped when the caller is only asking, because a declaration never mutates. + """ + held = self._grid_from(cache_attribute, [int(extent) for extent in shape]) + restored_shape = self._inverse_geometry(cache_attribute) + return held, self._grid_from(cache_attribute, restored_shape) 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" - ), - ) + if self.transforms is not None or self.displacement is not None: + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=( + "resampling through a map inverts to resampling through its inverse, and that" + " inverse is not declared here -- so a prediction finalize through this stage" + " assembles the volume. The forward direction streams" + ), + ) + try: + self._inverse_geometry(Attribute(cache_attribute)) + except NameError: + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=( + "the grid this stage resampled off is not on the attribute it is being asked to" + " invert, so the shape it restores is unknown here. The forward direction streams" + ), + ) + return PatchLocality(LocalityKind.REGRID) + + def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) -> list[int]: + try: + return self._inverse_geometry(Attribute(cache_attribute)) + except NameError: + return shape + + def inverse_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + del source_spatial_shape + self._inverse_geometry(cache_attribute) 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. 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 + if self._target_is_own and (self.transforms is not None or self.displacement is not None): + raise TransformError( + "'Resample' has no inverse here: it changes no grid, so undoing it is undoing its" + " map -- which is applying a different map, not this one backwards.", + "Set 'inverse: false' on this stage, or declare a second Resample with the inverse" + " transforms in the chain that needs it.", + ) + held, restored = self._inverse_grids(cache_attribute, [int(extent) for extent in tensor.shape[1:]]) + whole = tuple(slice(0, extent) for extent in restored.size_zyx) + return self._resample_between(restored, held, tensor, whole, [0] * restored.rank) + + def stream_region_inverse( + self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute + ) -> torch.Tensor: + del name + held, restored = self._inverse_grids(cache_attribute, [int(extent) for extent in context.source_shape]) + return self._resample_between( + restored, held, tensor, tuple(context.target), [part.start for part in context.source] ) def stream_region_target( - self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + self, name: str, 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.", - ) + del name + held, restored = self._inverse_grids(Attribute(cache_attribute), [int(e) for e in source_spatial_shape]) + identity = TransformBound.exact(AffineMap.identity(restored.rank)) + return list(source_window(restored.sub_grid(tuple(target_slices)), held, identity)) + def _resample_between( + self, + target: Grid, + source: Grid, + tensor: torch.Tensor, + target_slices: tuple[slice, ...], + region_starts: list[int], + ) -> torch.Tensor: + """One region of ``target``, read off ``source`` with no map between them.""" + region = target.sub_grid(target_slices) + shape, mode = list(source.size_zyx), self._mode(tensor) + axes = separable_source_index(region, source, (), tensor.device) + if axes is not None: + order = blend_order(target, source) + return gather_separable(tensor, axes, region_starts, shape, mode, self.fill_value, order) + coordinates = source_index(region, source, (), tensor.device) + return gather(tensor, coordinates, region_starts, shape, mode, self.fill_value) -class ResampleTransform(TransformInverse): - """Resample a volume through stored transforms (a displacement field, an affine). + @property + def _target_is_own(self) -> bool: + return isinstance(self._target, _OwnGrid) - Whole-volume: nothing in the format bounds the stored displacement, so no halo can be declared - from the header alone. - """ - def __init__(self, transforms: dict[str, bool], inverse: bool = True) -> None: - super().__init__(inverse) - self.transforms = transforms +class ResampleToResolution(Resample): + """Deprecated spelling of ``Resample: {spacing: ...}``.""" - def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: - return shape + def __init__(self, spacing: list[float] = [1.0, 1.0, 1.0], inverse: bool = True) -> None: + super().__init__(spacing=spacing, inverse=inverse) - def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - if len(tensor.shape) != 4: - raise NameError("Input size should be 5 dim") - _require_simpleitk() - image = data_to_image(tensor, cache_attribute) - transforms = [] - for transform_group, invert in self.transforms.items(): - transform = None - for dataset in self.datasets: - if dataset.is_dataset_exist(transform_group, name): - transform = dataset.read_transform(transform_group, name) - break - if transform is None: - raise NameError(f"Tranform : {transform_group}/{name} not found") - if isinstance(transform, sitk.BSplineTransform): - if invert: - transform_to_displacement_field_filter = sitk.TransformToDisplacementFieldFilter() - transform_to_displacement_field_filter.SetReferenceImage(image) - displacement_field = transform_to_displacement_field_filter.Execute(transform) - iterative_inverse_displacement_field_image_filter = ( - sitk.IterativeInverseDisplacementFieldImageFilter() - ) - iterative_inverse_displacement_field_image_filter.SetNumberOfIterations(20) - inverse_displacement_field = iterative_inverse_displacement_field_image_filter.Execute( - displacement_field - ) - transform = sitk.DisplacementFieldTransform(inverse_displacement_field) - else: - if invert: - transform = transform.GetInverse() - transforms.append(transform) - result_transform = sitk.CompositeTransform(transforms) - - # Resample through SimpleITK so the stored transform is applied in physical space: spacing, - # direction and the (x, y, z) mm units of the displacement are all honoured. A hand-rolled - # grid_sample would add the physical (dx, dy, dz) displacement straight onto a (z, y, x) - # voxel-index grid, transposing the x/z axes and treating millimetres as voxels. - interpolator = sitk.sitkNearestNeighbor if tensor.dtype == torch.uint8 else sitk.sitkLinear - resampled = sitk.Resample(image, image, result_transform, interpolator, 0.0) - data, _ = image_to_data(resampled) - result = torch.from_numpy(np.ascontiguousarray(data)) - return result.to(torch.uint8) if tensor.dtype == torch.uint8 else result.float() +class ResampleToShape(Resample): + """Deprecated spelling of ``Resample: {shape: ...}``.""" - def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - raise NotImplementedError( - "ResampleTransform.inverse is not implemented; set `inverse: false` on this transform " - "(it defaults to true)." + def __init__(self, shape: list[int] = [100, 256, 256], inverse: bool = True) -> None: + super().__init__(shape=shape, inverse=inverse) + + +class ResampleToReference(Resample): + """Deprecated spelling of ``Resample: {reference: ...}``.""" + + def __init__( + self, + 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, + interpolation: str | None = None, + inverse: bool = True, + ) -> None: + 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. Resample: {reference: 822174, reference_group: Volume}.", + ) + super().__init__( + reference=entry, + reference_group=group, + reference_dataset=dataset, + field=field, + field_group=field_group, + max_displacement=max_displacement, + fill=fill, + interpolation=interpolation, + inverse=inverse, ) +class ResampleTransform(Resample): + """Deprecated spelling of ``Resample: {transforms: ...}``.""" + + def __init__( + self, + transforms: dict[str, bool], + interpolation: str | None = None, + fill: float = 0.0, + inverse: bool = False, + ) -> None: + if not transforms: + raise TransformError( + "'ResampleTransform' needs at least one group of stored transforms to apply.", + "Name it and say whether to invert it, e.g. Resample: {transforms: {reg: false}}.", + ) + super().__init__(transforms=transforms, interpolation=interpolation, fill=fill, inverse=inverse) + + class Mask(Transform): """Set everything outside a mask to a constant. @@ -2614,32 +2388,8 @@ def _array_order_spacing(cache_attribute: Attribute) -> list[float] | 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. - """ +class Warp(Resample): + """Deprecated spelling of ``Resample: {field: ...}`` — a warp on the case's own grid.""" def __init__( self, @@ -2648,94 +2398,18 @@ def __init__( 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" - ), + "Declare it, e.g. Resample: {field: ./DVF:omezarr, max_displacement: 250.0}.", ) - 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. - - The field's components are (x, y, z) where the array axes are (z, y, x) -- the two orders - meet here, and reversing one of them is a warp that looks plausible and moves the anatomy - the wrong way. - """ - extent = list(tensor.shape[1:]) - # Grid and field follow the VOLUME's device: a field is read from disk onto the CPU, the - # volume may be GPU-resident, and grid_sample takes both from one device. - device = tensor.device - field = field.to(device) - axes = torch.meshgrid( - *[torch.arange(size, dtype=torch.float32, device=device) for size in extent], indexing="ij" - ) - sample = [] - for axis in range(len(extent)): - # field component for array axis `axis` (z,y,x) is the reversed one (x,y,z) - displacement = field[len(extent) - 1 - axis] / spacing[axis] - sample.append(axes[axis] + displacement) - # grid_sample wants the LAST dim ordered (x, y, z) and coordinates normalised to [-1, 1]. - grid = torch.stack( - [2.0 * sample[axis] / max(1, extent[axis] - 1) - 1.0 for axis in reversed(range(len(extent)))], dim=-1 + super().__init__( + field=field, + field_group=group, + max_displacement=max_displacement, + interpolation=interpolation, + inverse=False, ) - moved = F.grid_sample( - tensor.unsqueeze(0).float(), - grid.unsqueeze(0), - mode="bilinear" if self.interpolation == "linear" else "nearest", - padding_mode="zeros", - align_corners=True, - ) - return moved.squeeze(0).to(tensor.dtype) - - def stream_region( - self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute - ) -> torch.Tensor: - spacing = _array_order_spacing(cache_attribute) - if spacing is None: - return self(name, tensor, cache_attribute) - 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 = _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.displacement.read(name, None, len(tensor.shape) - 1) - # The declared bound is checked against every field read, on this path as on the streamed - # one: it is what sizes the halo, and a bound smaller than the field streams a region whose - # edge is missing. - self.displacement.check_bound(field, name) - return self._sample(tensor, field, spacing) class Reduce(Transform): @@ -2924,6 +2598,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -2944,6 +2619,7 @@ def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) def stream_region_target( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -2970,6 +2646,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -2987,12 +2664,13 @@ def stream_region_source( def stream_region_target( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, ) -> list[slice]: # A flip is its own inverse: a written region pulls exactly the region the forward would read. - return self.stream_region_source(target_slices, source_spatial_shape, cache_attribute) + return self.stream_region_source(name, target_slices, source_spatial_shape, cache_attribute) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: return tensor.flip(tuple(self.dims)) @@ -3142,6 +2820,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -3211,6 +2890,7 @@ def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) def stream_region_target( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, @@ -3680,6 +3360,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: def stream_region_source( self, + name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute, diff --git a/konfai/predictor.py b/konfai/predictor.py index 3ee127c8..6ae8f2b1 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -53,7 +53,6 @@ _halo_radii, _HaloPull, _RemapPull, - _ScalePull, blend_overlap, ) @@ -64,7 +63,7 @@ from konfai.data.transform import ( LocalityKind, PatchLocality, - Resample, + RegionContext, Transform, TransformInverse, TransformLoader, @@ -418,7 +417,7 @@ def __repr__(self) -> str: # The write-side region kinds: what SlabRegionStream can carry inside a streamed finalize chain (the # same set the read dispatcher accepts as its region stages; on both sides they compose). -_REGION_KINDS = (LocalityKind.HALO, LocalityKind.ORIENTATION, LocalityKind.CROP, LocalityKind.RESCALE) +_REGION_KINDS = (LocalityKind.HALO, LocalityKind.ORIENTATION, LocalityKind.CROP, LocalityKind.REGRID) # Streaming pays per-slab work (the pipe traversal, region writes, the TTA aligner); when the # assembled accumulators of all copies are below this fraction of allocatable memory (a 2.5D case), @@ -862,8 +861,7 @@ def _consume_slabs( ) -> None: """Run each jointly finalized slab through the plan: prefix per slab, then sink, region stream, or buffer. The first slab fixes the case's state (post-prefix attribute, region - scheduler or buffer) and may demote a RESCALE region to the buffered tail (see - ``_init_stream_state``).""" + scheduler or buffer; see ``_init_stream_state``).""" plan = cast(_StreamPlan, self._stream_plans[index]) for region, copies in slabs: block, attribute = self._finalize_slab(index, copies, number_of_channels_per_model, plan, dataset, region) @@ -887,11 +885,9 @@ def _init_stream_state( ) -> _StreamPlan: """Fix the case's streaming state at its first slab, when the prefix output is known. - A RESCALE stage streams through ``resample_region``, which matches ``F.interpolate`` bit for - bit in nearest mode (uint8) and to ~float-rounding in linear mode, so a rescale streams by - default and bounds a large float resample to a window. ``KONFAI_STREAM_LINEAR_RESAMPLE=0`` - demotes a float rescale here to the buffered whole-volume tail for a run that needs exactness; - the demotion leaves the prefix untouched (the pipe was never part of it). + A ``REGRID`` stage streams through the very code its whole-volume call runs, over a region + that happens to be smaller, so the streamed and whole-volume answers are equal by + construction rather than by agreement and a large resample never has to be held whole. """ self._post_prefix_attributes[index] = Attribute(attribute) spatial = [int(extent) for extent in self.output_layer_accumulator[index][0].shape] @@ -925,9 +921,7 @@ def _make_pipe_state( The fold is planned by walking a one-voxel corner of the real first slab through the pipe with one evolving attribute: each stage declares against, and remaps from, the state the stages before it left — a second resample pops the Size stack the first one already popped, - a reorientation after a permute reads the moved axes — and the walk carries the dtype, so a - float RESCALE (``resample_region`` matches ``F.interpolate`` only to ~float-rounding) answers - ``None`` — demoting to the whole-volume tail — only under ``KONFAI_STREAM_LINEAR_RESAMPLE=0``. + a reorientation after a permute reads the moved axes — and the walk carries the dtype. ``produce`` then replays the same transitions on a fresh copy per emission (the same slab-local scoping as the prefix). """ @@ -950,38 +944,28 @@ def _make_pipe_state( pull_fns.append(_HaloPull(_halo_radii(locality.halo, len(shape)), shape)) shapes.append(list(shape)) probe = stage(name, probe, walking) - elif locality.kind is LocalityKind.RESCALE: - # A float rescale streams within a window: resample_region computes the same linear - # taps the read side already streams, matching the whole-volume F.interpolate to - # ~float-rounding (a boundary voxel or two flips after argmax; a raw float output - # differs by ~1 ULP) -- which bounds a large float resample (a probability volume - # sent back to native) instead of holding it whole. KONFAI_STREAM_LINEAR_RESAMPLE=0 - # forces the exact whole-volume resample for a run that needs bit-identity. Nearest - # (uint8) is byte-identical either way. - if probe.dtype is not torch.uint8 and not env_flag("KONFAI_STREAM_LINEAR_RESAMPLE", True): - return None - resample = cast(Resample, stage.transform) + elif locality.kind is LocalityKind.REGRID: + # A regrid states its transition instead of performing it: its inverse restores a + # whole volume from a region, which is not something the one-voxel probe can be + # run through, and its forward answer here would be one voxel's target grid. + transform = stage.transform if stage.inverted: - pull_fns.append(_RemapPull(resample.stream_region_target, shape, snapshot)) - out = resample._inverse_geometry(walking) + remapper = cast(TransformInverse, transform) + pull_fns.append(_RemapPull(remapper.stream_region_target, shape, snapshot, name)) + out = remapper.inverse_transform_shape(list(shape), Attribute(walking)) + remapper.inverse_stream_cache_attribute(walking, shape) else: - out = [ - int(extent) - for extent in resample.transform_shape( - self.group_src, name, list(shape), Attribute(walking) - ) - ] - scales = [shape[k] / out[k] for k in range(len(shape))] - pull_fns.append(_ScalePull(scales, shape)) - resample.write_stream_cache_attribute(walking, shape) + pull_fns.append(_RemapPull(transform.stream_region_source, shape, snapshot, name)) + out = transform.transform_shape(self.group_src, name, list(shape), Attribute(walking)) + transform.write_stream_cache_attribute(walking, shape) shapes.append([int(extent) for extent in out]) elif locality.kind in _REGION_KINDS: if stage.inverted: remapper = cast(TransformInverse, stage.transform) - pull_fns.append(_RemapPull(remapper.stream_region_target, shape, snapshot)) + pull_fns.append(_RemapPull(remapper.stream_region_target, shape, snapshot, name)) out = remapper.inverse_transform_shape(list(shape), Attribute(walking)) else: - pull_fns.append(_RemapPull(stage.transform.stream_region_source, shape, snapshot)) + pull_fns.append(_RemapPull(stage.transform.stream_region_source, shape, snapshot, name)) out = stage.transform.transform_shape(self.group_src, name, list(shape), Attribute(walking)) shapes.append([int(extent) for extent in out]) # The stage's attribute transition, on a one-voxel corner: a crop's tensor answer @@ -1042,19 +1026,22 @@ def _apply_pipe_stage( # is one window's, dropped. stage(name, block, attribute) return block - if kind is LocalityKind.RESCALE: - resample = cast(Resample, stage.transform) - scales = [in_shape[k] / out_shape[k] for k in range(len(in_shape))] - result = resample.resample_region(block, target, [s.start for s in source], scales, in_shape) + if kind is LocalityKind.REGRID: + # Region-aware on both sides, and the geometry written from the FULL shape: what the + # stage records on the way is one region's, and the case's answer is the whole grid's. + context = RegionContext(tuple(source), tuple(target), tuple(in_shape), tuple(out_shape)) if stage.inverted: - resample._inverse_geometry(attribute) + remapper = cast(TransformInverse, stage.transform) + result = remapper.stream_region_inverse(name, block, context, Attribute(attribute)) + remapper.inverse_stream_cache_attribute(attribute, in_shape) else: - resample.write_stream_cache_attribute(attribute, in_shape) + result = stage.transform.stream_region(name, block, context, Attribute(attribute)) + stage.transform.write_stream_cache_attribute(attribute, in_shape) return result if kind is LocalityKind.ORIENTATION and not stage.inverted: # A forward orientation writes the case origin/direction from the extent it is handed; run # the tensor action on a throwaway scope so it does not record the SLAB's extent, then write - # the case geometry from the full ``in_shape`` (its documented contract) -- as RESCALE does. + # the case geometry from the full ``in_shape`` (its documented contract) -- as REGRID does. result = stage(name, block, Attribute(attribute)) cast(TransformInverse, stage.transform).write_stream_cache_attribute(attribute, in_shape) return result diff --git a/konfai/utils/ITK.py b/konfai/utils/ITK.py index b6b94b63..d30517fc 100644 --- a/konfai/utils/ITK.py +++ b/konfai/utils/ITK.py @@ -19,6 +19,7 @@ from __future__ import annotations from pathlib import Path +from typing import TYPE_CHECKING, cast import numpy as np import torch @@ -31,6 +32,9 @@ from konfai.utils.errors import TransformError +if TYPE_CHECKING: + from konfai.data.geometry import AffineMap, AffineStage, DisplacementStage, Grid, SpatialStages + def _require_simpleitk() -> None: """Raise a clear project error when an ITK-only path is used without SimpleITK.""" @@ -444,3 +448,117 @@ def clip_and_cast(image: sitk.Image, min_value: float, max_value: float, dtype: result = sitk.GetImageFromArray(data.astype(dtype)) result.CopyInformation(image) return result + + +# ------------------------------------------------------------------ decoding a stored transform +# The bridge from a sitk.Transform to konfai.data.geometry's sitk-free stages: everything a +# resample needs to sample and bound the map, as plain numpy, decoded once per case. + + +def _linear_map(transform: sitk.Transform) -> AffineMap: + """The exact world map of a linear transform: ``T(p) = M p + T(0)``. + + ``M`` comes from ``GetMatrix`` where the type has one -- the number ITK itself resamples with, + read past the centre/translation parameterisation that differs between Euler, Similarity, + Scale and Affine -- and from ``T(e_k) - T(0)`` otherwise. The offset is ``T(0)`` directly + rather than assembled from centre and translation: one call, no cancellation, and true for + every parameterisation at once. + + Probing is sound HERE and nowhere else in this file: for an affine map the columns are the map, + exactly, by linearity. For a non-linear one the same arithmetic measures a local gradient and + extrapolates it, which under-bounds -- which is why a BSpline's affine part is the identity and + all of its reach lives in the residual. + """ + from konfai.data.geometry import AffineMap + + rank = int(transform.GetDimension()) + offset = np.asarray(transform.TransformPoint((0.0,) * rank), dtype=np.float64) + if hasattr(transform, "GetMatrix"): + matrix = np.asarray(transform.GetMatrix(), dtype=np.float64).reshape(rank, rank) + else: + basis = np.eye(rank) + columns = [ + np.asarray(transform.TransformPoint(tuple(basis[k])), dtype=np.float64) - offset for k in range(rank) + ] + matrix = np.stack(columns, axis=1) + return AffineMap(matrix, offset) + + +def _grid_of_image(image: sitk.Image) -> Grid: + from konfai.data.geometry import Grid + + rank = int(image.GetDimension()) + return Grid( + tuple(int(extent) for extent in reversed(image.GetSize())), + np.asarray(image.GetOrigin(), dtype=np.float64), + np.asarray(image.GetSpacing(), dtype=np.float64), + np.asarray(image.GetDirection(), dtype=np.float64).reshape(rank, rank), + ) + + +def _displacement_stage(grid: Grid, per_component: list[np.ndarray], order: int, what: str) -> DisplacementStage: + from konfai.data.geometry import DisplacementStage + + values = np.stack([component.astype(np.float64, copy=False) for component in per_component]) + if not np.isfinite(values).all(): + raise TransformError( + f"{what} carries a non-finite displacement value, so no bound on its reach exists.", + "A NaN or infinite coefficient means the transform was written from a failed solve;" + " re-export it, or drop it from 'transforms:'.", + ) + return DisplacementStage(grid, values, order) + + +def decode_transform_stages(transform: sitk.Transform) -> SpatialStages: + """A stored transform as geometry stages in APPLICATION order, or a refusal naming the type. + + ``CompositeTransform`` applies its member list in REVERSE (the last added runs first — verified + against SimpleITK, where ``GetNthTransform(0)`` is nonetheless the first added); the reversal is + normalized here, once, so every consumer reads stages first-applied-first. + """ + _require_simpleitk() + if isinstance(transform, sitk.CompositeTransform): + stages: list[AffineStage | DisplacementStage] = [] + for index in reversed(range(transform.GetNumberOfTransforms())): + stages.extend(decode_transform_stages(transform.GetNthTransform(index))) + return tuple(stages) + from konfai.data.geometry import AffineStage + + if isinstance(transform, sitk.BSplineTransform): + coefficients = transform.GetCoefficientImages() + arrays = [sitk.GetArrayFromImage(component) for component in coefficients] + return ( + _displacement_stage( + _grid_of_image(coefficients[0]), arrays, int(transform.GetOrder()), "this BSpline transform" + ), + ) + if isinstance(transform, sitk.DisplacementFieldTransform): + field = transform.GetDisplacementField() + array = sitk.GetArrayFromImage(field) # (Z, Y, X, rank), components (x, y, z) + components = [np.ascontiguousarray(array[..., k]) for k in range(array.shape[-1])] + return (_displacement_stage(_grid_of_image(field), components, 1, "this displacement field"),) + if transform.IsLinear(): + return (AffineStage(_linear_map(transform)),) + raise TransformError( + f"A stored '{transform.GetName()}' decomposes into no bounded map: how far a target region" + " reaches into its source is unknown, so the region it must read is unbounded.", + "Convert it to a displacement field when it is written, or use a rigid/affine/BSpline" + " transform, which all decompose.", + ) + + +def invert_stages(stages: SpatialStages, rank: int) -> SpatialStages | None: + """The exact inverse of an all-affine decoded map, or ``None`` when one is not algebraic. + + A BSpline or a field inverts by an iterative dense solve, not an algebraic step, and a field + solved per region is not the restriction of the field solved once — so a non-affine inverse is + ``None`` here and the caller refuses with the remedy, rather than resampling through a guess. + """ + from konfai.data.geometry import AffineMap, AffineStage + + if not all(isinstance(stage, AffineStage) for stage in stages): + return None + folded = AffineMap.identity(rank) + for stage in stages: + folded = folded.then(cast("AffineStage", stage).map) + return (AffineStage(folded.inverted()),) diff --git a/tests/integration/test_transform_doc_examples.py b/tests/integration/test_transform_doc_examples.py index a3ab1283..840a62e3 100644 --- a/tests/integration/test_transform_doc_examples.py +++ b/tests/integration/test_transform_doc_examples.py @@ -141,6 +141,23 @@ def _write_fields(root: Path, cases: int = 2, shape=(4, 6, 6)) -> None: sitk.WriteImage(image, str(case / "DVF.mha")) +def _write_transforms(root: Path, cases: int = 2) -> None: + """One stored transform per case, beside the images, as a group of its own. + + A page that documents applying a registration solved elsewhere has to be able to show one, and + an example is only runnable if the fixture holds what it names. KonfAI reads a stored transform + from ``/.itk.txt``, which is what SimpleITK writes. + """ + for index in range(cases): + case = root / f"case_{index}" + case.mkdir(parents=True, exist_ok=True) + transform = sitk.Euler3DTransform() + transform.SetCenter((3.0, 5.0, 11.0)) + transform.SetRotation(0.05, -0.03, 0.08) + transform.SetTranslation((0.4, -0.6, 1.1)) + sitk.WriteTransform(transform, str(case / "reg.itk.txt")) + + def _doc_examples() -> list[tuple[int, str]]: # Asserted, not skipped: an empty parametrize set is a pass, so a page that moved would retire # this whole guard without a single red test. @@ -156,6 +173,7 @@ def test_a_documented_example_plans_and_runs(line: int, config: str, tmp_path: P workdir.mkdir() _write_dataset(workdir / "Raw", _source_groups(config)) _write_fields(workdir / "Fields") + _write_transforms(workdir / "Raw") (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_geometry.py b/tests/unit/test_geometry.py new file mode 100644 index 00000000..999c59a6 --- /dev/null +++ b/tests/unit/test_geometry.py @@ -0,0 +1,178 @@ +# 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 + +"""The geometry vocabulary against SimpleITK, not against itself. + +Every rule here is pinned to the external oracle: two KonfAI paths agree by construction, +including on a grid placed in the wrong place. +""" + +import itertools + +import numpy as np +import pytest +from konfai.data.geometry import AffineMap, Grid, TransformBound, WorldBox +from konfai.utils.dataset import Attribute +from konfai.utils.errors import TransformError + +sitk = pytest.importorskip("SimpleITK") + + +def _oblique_direction() -> np.ndarray: + a, b, c = np.deg2rad(17.0), np.deg2rad(-23.0), np.deg2rad(11.0) + rz = np.array([[np.cos(a), -np.sin(a), 0.0], [np.sin(a), np.cos(a), 0.0], [0.0, 0.0, 1.0]]) + ry = np.array([[np.cos(b), 0.0, np.sin(b)], [0.0, 1.0, 0.0], [-np.sin(b), 0.0, np.cos(b)]]) + rx = np.array([[1.0, 0.0, 0.0], [0.0, np.cos(c), -np.sin(c)], [0.0, np.sin(c), np.cos(c)]]) + return rz @ ry @ rx + + +def _grid(direction: np.ndarray | None = None) -> Grid: + return Grid( + size_zyx=(29, 33, 48), + origin_xyz=np.array([-31.0, 12.5, 4.25]), + spacing_xyz=np.array([0.7, 1.3, 0.9]), + direction_xyz=np.eye(3) if direction is None else direction, + ) + + +def _image(grid: Grid) -> "sitk.Image": + image = sitk.GetImageFromArray(np.zeros(grid.size_zyx, np.float32)) + image.SetOrigin(tuple(grid.origin_xyz)) + image.SetSpacing(tuple(grid.spacing_xyz)) + image.SetDirection(tuple(grid.direction_xyz.ravel())) + return image + + +class TestGridAgainstSimpleITK: + def test_index_to_world_is_transform_index_to_physical_point(self): + grid = _grid(_oblique_direction()) + image = _image(grid) + for index_xyz in ([0.0, 0.0, 0.0], [47.0, 32.0, 28.0], [3.5, -0.5, 17.25]): + want = np.array(image.TransformContinuousIndexToPhysicalPoint(index_xyz)) + got = grid.index_to_world.apply(np.asarray(index_xyz)) + np.testing.assert_allclose(got, want, rtol=0.0, atol=1e-12) + + def test_sub_grid_origin_is_the_slab_origin_bit_for_bit(self): + # The load-bearing line: same product, same association as ITK, on an oblique grid with + # anisotropic spacing and a non-round origin. + grid = _grid(_oblique_direction()) + image = _image(grid) + region = (slice(11, 21), slice(0, 33), slice(5, 48)) + want = np.array(image.TransformIndexToPhysicalPoint([5, 0, 11])) + sub = grid.sub_grid(region) + assert sub.size_zyx == (10, 33, 43) + np.testing.assert_array_equal(sub.origin_xyz, want) + + def test_world_to_index_round_trips(self): + grid = _grid(_oblique_direction()) + points = np.random.RandomState(0).uniform(-200, 200, size=(50, 3)) + back = grid.world_to_index.apply(grid.index_to_world.apply(points)) + np.testing.assert_allclose(back, points, rtol=0.0, atol=1e-9) + + def test_world_box_covers_the_outer_faces(self): + grid = _grid(_oblique_direction()) + image = _image(grid) + box = grid.world_box((slice(4, 9), slice(0, 33), slice(0, 48))) + corners = itertools.product((3.5, 8.5), (-0.5, 32.5), (-0.5, 47.5)) + for z, y, x in corners: + point = np.array(image.TransformContinuousIndexToPhysicalPoint([x, y, z])) + assert np.all(point >= box.low_xyz - 1e-9) and np.all(point <= box.high_xyz + 1e-9) + + +class TestGridOf: + def _attribute(self, **overrides: object) -> Attribute: + attribute = Attribute() + values: dict[str, np.ndarray] = { + "Origin": np.array([-31.0, 12.5, 4.25]), + "Spacing": np.array([0.7, 1.3, 0.9]), + "Direction": np.eye(3).ravel(), + } + values.update({key: np.asarray(value) for key, value in overrides.items()}) + for key, value in values.items(): + attribute[key] = value + return attribute + + def test_reads_a_full_header(self): + grid = Grid.of([29, 33, 48], self._attribute(), "case 'A'") + assert grid.size_zyx == (29, 33, 48) + np.testing.assert_array_equal(grid.spacing_xyz, [0.7, 1.3, 0.9]) + + def test_refuses_a_missing_key_naming_it(self): + attribute = self._attribute() + attribute.pop("Spacing") + with pytest.raises(TransformError, match="Spacing"): + Grid.of([29, 33, 48], attribute, "case 'A'") + + def test_refuses_a_non_positive_spacing(self): + with pytest.raises(TransformError, match="positive"): + Grid.of([29, 33, 48], self._attribute(Spacing=np.array([0.7, 0.0, 0.9])), "case 'A'") + + def test_readable_is_total_and_quiet(self): + assert Grid.readable(self._attribute()) + assert not Grid.readable(Attribute()) + + +class TestWorldBoxImage: + def test_image_under_equals_the_corner_hull(self): + # The identity that licenses never enumerating corners: |A| on the half-extents equals the + # hull of the 2^rank mapped corners, for any affine. + rng = np.random.RandomState(7) + for _ in range(200): + affine = AffineMap(rng.randn(3, 3), rng.randn(3) * 100.0) + low = rng.uniform(-100, 0, 3) + high = low + rng.uniform(0.1, 200, 3) + box = WorldBox(low, high) + corners = np.array([affine.apply(np.array(c)) for c in itertools.product(*zip(low, high, strict=True))]) + image = box.image_under(affine) + np.testing.assert_allclose(image.low_xyz, corners.min(axis=0), rtol=0.0, atol=1e-9) + np.testing.assert_allclose(image.high_xyz, corners.max(axis=0), rtol=0.0, atol=1e-9) + + +class TestAffineMap: + def test_then_composes_in_stated_order(self): + inner = AffineMap(np.diag([2.0, 1.0, 1.0]), np.array([1.0, 0.0, 0.0])) + outer = AffineMap(np.eye(3), np.array([0.0, 10.0, 0.0])) + point = np.array([1.0, 1.0, 1.0]) + np.testing.assert_array_equal(inner.then(outer).apply(point), outer.apply(inner.apply(point))) + + def test_inverted_refuses_a_singular_matrix(self): + singular = AffineMap(np.diag([1.0, 0.0, 1.0]), np.zeros(3)) + with pytest.raises(TransformError, match="singular"): + singular.inverted() + + +class TestTransformBound: + def test_after_transports_the_residual_through_the_outer_matrix(self): + # T = A ∘ (id + d): the residual scales by |A|, not by 1. An affine of scale 3 around a + # 5 mm spline reaches 15 mm — the naive bound is short by 10 voxels of every slab edge. + scale = TransformBound.exact(AffineMap(np.diag([3.0, 1.0, 1.0]), np.zeros(3))) + wiggle = TransformBound.shift(np.array([5.0, 5.0, 5.0])) + composed = scale.after(wiggle) + np.testing.assert_array_equal(composed.residual_xyz, [15.0, 5.0, 5.0]) + # And the other order leaves it alone. + np.testing.assert_array_equal(wiggle.after(scale).residual_xyz, [5.0, 5.0, 5.0]) + + def test_map_box_contains_the_true_image_of_a_nonlinear_map(self): + rng = np.random.RandomState(3) + affine = AffineMap(np.eye(3) + rng.randn(3, 3) * 0.1, rng.randn(3) * 10.0) + residual = np.array([4.0, 2.0, 1.0]) + bound = TransformBound(affine, residual) + box = WorldBox(np.array([-20.0, -10.0, 0.0]), np.array([15.0, 25.0, 30.0])) + mapped = bound.map_box(box) + points = rng.uniform(box.low_xyz, box.high_xyz, size=(500, 3)) + # Any map of the form affine + bounded wiggle stays inside. + images = affine.apply(points) + rng.uniform(-1.0, 1.0, size=(500, 3)) * residual + assert np.all(images >= mapped.low_xyz - 1e-9) and np.all(images <= mapped.high_xyz + 1e-9) diff --git a/tests/unit/test_resample.py b/tests/unit/test_resample.py new file mode 100644 index 00000000..80ccb135 --- /dev/null +++ b/tests/unit/test_resample.py @@ -0,0 +1,383 @@ +# 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 + +"""One ``Resample``: which grid to write on, what map to write it through, and what that fixed. + +``ResampleToResolution``, ``ResampleToShape``, ``ResampleToReference``, ``ResampleTransform`` and +``Warp`` were five stages answering two questions between them, each with its own sampler and its +own idea of where a voxel is. They are now five spellings of this one, and the tests here are for +what only became checkable once there was a single answer to check. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from konfai.data.transform import ( + Resample, + ResampleToReference, + ResampleToResolution, + ResampleToShape, + ResampleTransform, + Warp, +) +from konfai.utils.dataset import Attribute +from konfai.utils.errors import TransformError + +sitk = pytest.importorskip("SimpleITK") + +_ORIGIN, _SPACING = [11.5, -3.25, 40.0], [1.3, 0.9, 0.9] +_SHAPE = (24, 30, 32) + + +def _attributes(origin=None, spacing=None, direction=None) -> Attribute: + attribute = Attribute() + attribute["Origin"] = np.asarray(_ORIGIN if origin is None else origin, dtype=np.float64) + attribute["Spacing"] = np.asarray(_SPACING if spacing is None else spacing, dtype=np.float64) + attribute["Direction"] = (np.eye(3) if direction is None else direction).reshape(-1) + return attribute + + +def _volume(shape=_SHAPE, seed: int = 0) -> np.ndarray: + # Noise, not a smooth field: a smooth volume resampled onto a grid that is half a voxel off is + # still nearly right, so a smooth fixture would pass a map that is wrong by exactly the amount + # this file exists to catch. + return np.random.default_rng(seed).normal(size=shape).astype(np.float32)[None] + + +def _as_image(volume: np.ndarray, attribute: Attribute) -> sitk.Image: + image = sitk.GetImageFromArray(volume[0]) + image.SetOrigin(attribute.get_np_array("Origin").tolist()) + image.SetSpacing(attribute.get_np_array("Spacing").tolist()) + image.SetDirection(attribute.get_np_array("Direction").tolist()) + return image + + +# ------------------------------------------------------------------ one class, five spellings + + +@pytest.mark.parametrize( + ("alias", "unified"), + [ + (lambda: ResampleToResolution(spacing=[2.0, 1.5, 1.5]), lambda: Resample(spacing=[2.0, 1.5, 1.5])), + (lambda: ResampleToShape(shape=[12, 20, 22]), lambda: Resample(shape=[12, 20, 22])), + (lambda: ResampleToShape(shape=[0, 20, 0]), lambda: Resample(shape=[0, 20, 0])), + ], +) +def test_a_spelling_and_the_unified_stage_are_the_same_stage(alias, unified) -> None: + """The published names are argument translations, not behaviour of their own. + + Kept because they appear in shipped configs and in every bundle on the hub -- and kept THIN, + because a spelling that carries logic is a second implementation waiting to drift. + """ + volume = torch.from_numpy(_volume()) + left, right = alias(), unified() + assert isinstance(left, Resample) + assert left.apply_inverse == right.apply_inverse + + got = left("case", volume.clone(), _attributes()) + want = right("case", volume.clone(), _attributes()) + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +def test_every_spelling_is_the_one_class() -> None: + for stage in ( + ResampleToResolution(), + ResampleToShape(), + ResampleToReference(entry="x"), + ResampleTransform(transforms={"reg": False}), + Warp(field="./x:h5", group="DVF"), + ): + assert isinstance(stage, Resample) + + +def test_the_three_ways_to_name_a_target_grid_are_exclusive() -> None: + with pytest.raises(TransformError, match="three ways to say the same thing"): + Resample(spacing=[1.0, 1.0, 1.0], shape=[4, 4, 4]) + + +# ------------------------------------------------------------------ where the new grid sits + + +def test_extent_alignment_keeps_the_field_of_view_and_origin_alignment_keeps_voxel_zero() -> None: + """The one silent choice in the family, made explicit. + + ``extent`` makes the outer faces coincide, which is ``F.interpolate``'s map and what KonfAI has + always done; ``origin`` keeps voxel zero's centre where it is, which is what resampling onto a + grid that shares an origin does. A quarter of a voxel of anatomy separates them. + """ + attribute = _attributes() + volume = torch.from_numpy(_volume()) + + Resample(spacing=[2.0, 1.5, 1.5], align="extent")("case", volume.clone(), attribute) + extent_origin = attribute.get_np_array("Origin").copy() + extent_spacing = attribute.get_np_array("Spacing").copy() + + attribute = _attributes() + Resample(spacing=[2.0, 1.5, 1.5], align="origin")("case", volume.clone(), attribute) + + np.testing.assert_allclose(attribute.get_np_array("Origin"), _ORIGIN) + np.testing.assert_allclose(attribute.get_np_array("Spacing"), [2.0, 1.5, 1.5]) + # Extent alignment puts voxel zero half the spacing change away, on every axis. + np.testing.assert_allclose(extent_origin, np.asarray(_ORIGIN) + 0.5 * (extent_spacing - np.asarray(_SPACING))) + + +def test_an_unknown_alignment_is_refused_at_construction() -> None: + with pytest.raises(TransformError, match="unknown align"): + Resample(spacing=[1.0, 1.0, 1.0], align="corners") + + +# ------------------------------------------------------------------ the header describes the data + + +@pytest.mark.parametrize( + "kwargs", + [ + {"spacing": [2.0, 1.5, 1.5]}, + {"spacing": [2.0, 1.5, 1.5], "align": "origin"}, + {"shape": [12, 20, 22]}, + {"spacing": [0.9, 1.7, 1.1]}, + ], +) +def test_the_recorded_header_describes_the_grid_that_was_actually_sampled(kwargs) -> None: + """The check neither predecessor could pass, because neither wrote a placement at all. + + ``ResampleToResolution`` recorded the spacing that was ASKED FOR while sampling at ``n_in/n_out`` + times the source's -- up to a millimetre of drift across a volume -- and left the Origin alone + while sampling half a spacing-change away from it. Nothing downstream could see either: the + voxels are all real, and the header is the only witness. + + So the oracle is built FROM THE RECORDED HEADER. If the two disagree, resampling the source onto + the grid the header claims cannot reproduce the voxels the stage returned. + """ + volume = _volume() + attribute = _attributes() + got = Resample(**kwargs)("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] + + grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) + grid.SetOrigin(attribute.get_np_array("Origin").tolist()) + grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) + grid.SetDirection(attribute.get_np_array("Direction").tolist()) + want = sitk.GetArrayFromImage( + sitk.Resample(_as_image(volume, _attributes()), grid, sitk.Transform(), sitk.sitkLinear, 0.0) + ) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) + + +def test_an_oblique_case_records_an_oblique_header() -> None: + """A direction is carried through, not quietly dropped -- and the data follows it.""" + angle = np.deg2rad(23.0) + cos, sin = float(np.cos(angle)), float(np.sin(angle)) + direction = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]) + volume = _volume() + attribute = _attributes(direction=direction) + got = Resample(spacing=[2.0, 1.5, 1.5])("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] + + np.testing.assert_allclose(attribute.get_np_array("Direction").reshape(3, 3), direction) + grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) + grid.SetOrigin(attribute.get_np_array("Origin").tolist()) + grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) + grid.SetDirection(direction.reshape(-1).tolist()) + want = sitk.GetArrayFromImage( + sitk.Resample(_as_image(volume, _attributes(direction=direction)), grid, sitk.Transform(), sitk.sitkLinear, 0.0) + ) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) + + +# ------------------------------------------------------------------ an image and its label map + + +def test_a_label_map_lands_on_the_same_voxels_as_the_image_beside_it() -> None: + """The bug that had no symptom: a mask resampled with its CT came out shifted against it. + + ``F.interpolate``'s nearest reads ``floor(o * scale)`` where its linear reads + ``scale * (o + 0.5) - 0.5`` -- so the label map lagged the image of the SAME stage by + ``(scale - 1) / 2`` source voxels. At 0.5 mm resampled to 3 mm that is 2.5 source voxels, 1.25 mm + of anatomy, and both volumes are entirely plausible on their own. + + A ramp makes it visible: linear interpolation of a linear function is exact, so the resampled + image IS the continuous source coordinate, and the resampled label map must be its rounding. + """ + extent = 48 + ramp = np.broadcast_to(np.arange(extent, dtype=np.float32).reshape(1, 1, -1), (1, 4, 4, extent)) + attribute = _attributes(origin=[0.0, 0.0, 0.0], spacing=[0.5, 1.0, 1.0]) + + image = Resample(spacing=[3.0, 1.0, 1.0])("case", torch.from_numpy(np.ascontiguousarray(ramp)), attribute) + labels = Resample(spacing=[3.0, 1.0, 1.0])( + "case", torch.from_numpy(np.ascontiguousarray(ramp).astype(np.uint8)), _attributes([0.0] * 3, [0.5, 1.0, 1.0]) + ) + + coordinate = image.numpy()[0, 0, 0] + picked = labels.numpy()[0, 0, 0].astype(np.int64) + # The edges clamp, so the interior is where the ramp still reads its own coordinate. + interior = slice(1, -1) + np.testing.assert_array_equal(picked[interior], np.floor(coordinate[interior] + 0.5).astype(np.int64)) + assert float(np.abs(picked[interior] - coordinate[interior]).max()) <= 0.5 + + +# ------------------------------------------------------------------ the count + + +def test_a_spacing_that_binary_cannot_hold_does_not_lose_a_slice() -> None: + """90 voxels of 0.7 mm re-cut at 1.5 mm is 42.0 -- and in float64 it is 41.999999999999997. + + Truncating that gives 41: one slice of anatomy dropped, and a recorded spacing that no longer + covers what was read. The old float32 round-trip happened to land above; nothing said so. + """ + attribute = _attributes(origin=[0.0] * 3, spacing=[0.7, 0.7, 0.7]) + shape = Resample(spacing=[1.5, 1.5, 1.5]).transform_shape("", "case", [90, 90, 90], attribute) + assert shape == [42, 42, 42] + + +# ------------------------------------------------------------------ one grid change, one map + + +def test_a_change_of_grid_and_a_warp_are_one_interpolation(tmp_path) -> None: + """Asked for together they compose into one coordinate per voxel, checked against sitk's own. + + Two stages would interpolate the same voxels twice, and the second pass invents none of what the + first smoothed away -- which is the whole reason an atlas's appearance is rebuilt from native + volumes rather than from warped ones. + """ + from konfai.utils.dataset import Dataset + + field_shape, field_origin, field_spacing = (6, 8, 9), [10.0, -5.0, 38.0], [4.0, 3.5, 3.0] + field = np.zeros((3, *field_shape), dtype=np.float32) + for component, value in enumerate((1.5, -2.0, 0.75)): + field[component] = value + + store = Dataset(tmp_path / "DVF", "h5") + store.write("DVF", "case", field, _attributes(field_origin, field_spacing)) + + volume = _volume() + attribute = _attributes() + stage = Resample(spacing=[2.0, 1.5, 1.5], field=str(tmp_path / "DVF") + ":h5", field_group="DVF") + stage.set_datasets([store]) + got = stage("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] + + grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) + grid.SetOrigin(attribute.get_np_array("Origin").tolist()) + grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) + vector = sitk.GetImageFromArray(np.moveaxis(field, 0, -1).astype(np.float64), isVector=True) + vector.SetOrigin(field_origin) + vector.SetSpacing(field_spacing) + want = sitk.GetArrayFromImage( + sitk.Resample( + _as_image(volume, _attributes()), + grid, + sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)), + sitk.sitkLinear, + 0.0, + ) + ) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-3) + + +def test_a_bound_declared_without_a_field_is_refused() -> None: + with pytest.raises(TransformError, match="no field to apply"): + Resample(spacing=[1.0, 1.0, 1.0], max_displacement=10.0) + + +# ------------------------------------------------------------------ which loop runs + + +def test_a_map_that_factorises_takes_the_separable_loop() -> None: + """The optimisation needs a test, or it can be lost to a refactor with everything still green. + + A grid change between axis-aligned volumes reads one axis at a time; a rotation between them, or + a displacement, cannot and falls to the coordinate volume. Both are correct — the difference is + 43x versus 660x of ``F.interpolate`` on a CT-sized case, which no assertion about values shows. + """ + from konfai.data.geometry import Grid + from konfai.data.sampling import separable_source_index + + device = torch.device("cpu") + source = Grid(_SHAPE, np.asarray(_ORIGIN), np.asarray(_SPACING), np.eye(3)) + aligned = source.resampled(spacing_xyz=np.asarray([2.0, 1.5, 1.5])) + assert separable_source_index(aligned, source, (), device) is not None + + angle = np.deg2rad(23.0) + cos, sin = float(np.cos(angle)), float(np.sin(angle)) + turned = Grid( + aligned.size_zyx, + aligned.origin_xyz, + aligned.spacing_xyz, + np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]), + ) + assert separable_source_index(turned, source, (), device) is None, "a rotation does not factorise" + + # A flip is still axis-aligned, so it factorises -- the test that the check is not merely + # "is the direction the identity". + flipped = Grid(source.size_zyx, source.origin_xyz, source.spacing_xyz, np.diag([-1.0, -1.0, 1.0])) + assert separable_source_index(aligned, flipped, (), device) is not None + + +def test_the_two_loops_agree_where_both_can_serve_the_same_map() -> None: + """Different summation orders, same answer to float rounding — and the same fill, exactly.""" + from konfai.data.geometry import Grid + from konfai.data.sampling import gather, gather_separable, separable_source_index, source_index + + device = torch.device("cpu") + volume = torch.from_numpy(_volume()) + source = Grid(_SHAPE, np.asarray(_ORIGIN), np.asarray(_SPACING), np.eye(3)) + # Placed so part of the target reaches past the case, which is where the fill rule shows. + target = Grid((14, 18, 20), np.asarray(_ORIGIN) - 4.0, np.asarray([2.0, 1.6, 1.6]), np.eye(3)) + + axes = separable_source_index(target, source, (), device) + assert axes is not None + fast = gather_separable(volume, axes, [0, 0, 0], list(_SHAPE), "linear", -999.0) + general = gather(volume, source_index(target, source, (), device), [0, 0, 0], list(_SHAPE), "linear", -999.0) + + np.testing.assert_array_equal((fast == -999.0).numpy(), (general == -999.0).numpy()) + torch.testing.assert_close(fast, general, rtol=1e-6, atol=1e-5) + + +def test_the_blend_order_is_the_same_for_a_region_as_for_the_whole_volume() -> None: + """Axes are blended most-reduced-first, and the key must not be the extents in hand. + + Blending an axis reduces it before the next one reads it, so the order decides how much data + every later pass moves -- 9x on a thick-slice CT brought to isotropic. But the order also decides + the SUMMATION order, so a region that chose differently from the whole volume would stop being + bit-identical to it, which is the one equality the streaming design rests on. Keyed on the two + grids' spacings, which a region shares with its volume, and not on their extents, which it does + not. + """ + from konfai.data.geometry import Grid + from konfai.data.sampling import blend_order + + source = Grid((64, 512, 512), np.zeros(3), np.asarray([0.7, 0.7, 3.0]), np.eye(3)) + target = source.resampled(spacing_xyz=np.asarray([1.0, 1.0, 1.0])) + whole = blend_order(target, source) + + # z triples while y and x shrink, so y and x are blended first. + assert whole == [1, 2, 0] + for start, stop in ((0, 8), (17, 41), (target.size_zyx[0] - 3, target.size_zyx[0])): + region = target.sub_grid((slice(start, stop), slice(0, target.size_zyx[1]), slice(0, target.size_zyx[2]))) + assert blend_order(region, source) == whole + + +def test_an_axis_the_map_leaves_alone_is_left_alone() -> None: + """A resample of one axis reads the other two, it does not blend them — and says so in the values. + + ``spacing: [-1, -1, 3]`` keeps x and y exactly. Blending them anyway would be two gathers and a + lerp over the largest tensor in flight, for a result equal to the input; skipping them has to be + exactly that, not nearly. + """ + attribute = _attributes(origin=[0.0] * 3, spacing=[1.0, 1.0, 1.0]) + volume = torch.from_numpy(_volume()) + kept = Resample(spacing=[-1.0, -1.0, 1.0])("case", volume.clone(), Attribute(attribute)) + torch.testing.assert_close(kept, volume, rtol=0, atol=0) diff --git a/tests/unit/test_resample_sampler_rules.py b/tests/unit/test_resample_sampler_rules.py index 71c0c33e..b36042b2 100644 --- a/tests/unit/test_resample_sampler_rules.py +++ b/tests/unit/test_resample_sampler_rules.py @@ -14,17 +14,15 @@ # # SPDX-License-Identifier: Apache-2.0 -"""The rules every sampler in this package obeys, pinned separately from any one sampler. +"""The rules the one gather obeys, pinned apart from any stage that uses it. -There are two gather strategies for one arithmetic. ``Resample._resample_offset_region`` maps each -axis independently, so no coordinate volume is built and one ``index_select`` per axis does the work; -``ResampleToReference._sample_at`` cannot, because a displacement is not separable, so it holds a -coordinate per voxel and gathers eight corners flat. Same rules, different loops, and the loops are -different for a measured reason. +There is a single sampler in KonfAI — ``konfai.data.sampling.gather`` — and every resample, warp and +regrid reaches its voxels through it. That is recent: the rules used to be restated by a separable +sampler and a non-separable one, which is how two of them came to disagree about a half-voxel rim. -Rules kept apart from loops is only true while something checks it. These tests are that: they assert -the RULES -- the inside interval, the tap clamp, round-half-up, the working dtype, the fill -- against -SimpleITK and against each other, so a change to one gather cannot quietly stop matching the other. +One implementation does not make the rules self-evident, it only makes them checkable in one place. +These tests are that place: the inside interval, the tap clamp, round-half-up, the working dtype and +the fill, asserted against SimpleITK and against hand arithmetic. """ from __future__ import annotations @@ -32,38 +30,41 @@ import numpy as np import pytest import torch -from konfai.data.transform import Resample - - -class _Sampler(Resample): - """A bare handle on the sampler: these rules belong to `Resample`, not to any stage using it.""" - - def __call__(self, name, tensor, cache_attribute): # pragma: no cover - not the surface tested - raise NotImplementedError - - def write_stream_cache_attribute(self, cache_attribute, source_spatial_shape) -> None: - raise NotImplementedError - - def transform_shape(self, shape, cache_attribute): # pragma: no cover - raise NotImplementedError - - def inverse(self, name, tensor, cache_attribute): # pragma: no cover - raise NotImplementedError - - def patch_locality(self, cache_attribute): # pragma: no cover - raise NotImplementedError - +from konfai.data.sampling import gather _SOURCE = (12, 14, 16) _SCALES = [1.31, 1.17, 1.23] _OFFSETS = [0.4, -0.3, 0.2] -_TARGET = (slice(0, 8), slice(0, 9), slice(0, 10)) +_TARGET = (8, 9, 10) + +def _coordinates( + target_shape: tuple[int, ...] = _TARGET, + scales: list[float] | None = None, + offsets: list[float] | None = None, +) -> torch.Tensor: + """One source index per target voxel for the separable map ``scale * o + offset``, per ARRAY axis. -def _sampler(fill: float = 0.0) -> _Sampler: - sampler = _Sampler(inverse=False) - sampler.fill_value = fill - return sampler + Separable is the easy case to write by hand, not a second code path: what the gather receives is + always a coordinate per voxel, and how it was produced is none of its business. + """ + scales = _SCALES if scales is None else scales + offsets = _OFFSETS if offsets is None else offsets + axes = [ + scales[axis] * torch.arange(extent, dtype=torch.float64) + offsets[axis] + for axis, extent in enumerate(target_shape) + ] + grids = torch.meshgrid(*axes, indexing="ij") + # The gather wants the physical components last, in (x, y, z) — the mirror of the array axes. + return torch.stack(list(reversed(grids)), dim=-1) + + +def _sample(tensor: torch.Tensor, fill: float = 0.0, mode: str | None = None, **overrides) -> torch.Tensor: + source_shape = list(overrides.pop("source_shape", _SOURCE)) + coordinates = _coordinates(**overrides) + if mode is None: + mode = "nearest" if tensor.dtype == torch.uint8 else "linear" + return gather(tensor, coordinates, [0] * len(source_shape), source_shape, mode, fill) def _volume(offset: float = 0.0) -> np.ndarray: @@ -71,18 +72,6 @@ def _volume(offset: float = 0.0) -> np.ndarray: return (rng.random((1, *_SOURCE)) * 400 + offset).astype(np.float32) -def _offset_region(tensor: torch.Tensor, fill: float = 0.0, **overrides) -> torch.Tensor: - arguments = { - "target_slices": _TARGET, - "region_starts": [0, 0, 0], - "scales": _SCALES, - "n_in": list(_SOURCE), - "offsets": _OFFSETS, - } - arguments.update(overrides) - return _sampler(fill)._resample_offset_region(tensor, **arguments) # type: ignore[arg-type] - - def test_the_working_dtype_rule_holds_for_a_cpu_half_volume() -> None: """A CPU half is accumulated in float32; the eight-corner sum is not done in half. @@ -91,8 +80,8 @@ def test_the_working_dtype_rule_holds_for_a_cpu_half_volume() -> None: the two: accumulating in half drifts more than twice as far from the float32 answer. """ volume = _volume(offset=2050.0) # 2050..2450, entirely above 2048, where float16 spacing is 2 - guarded = _offset_region(torch.from_numpy(volume).half()).float() - reference = _offset_region(torch.from_numpy(volume)) + guarded = _sample(torch.from_numpy(volume).half()).float() + reference = _sample(torch.from_numpy(volume)) assert guarded.dtype is torch.float32 drift = float((guarded - reference).abs().max()) @@ -103,18 +92,20 @@ def test_the_working_dtype_rule_holds_for_a_cpu_half_volume() -> None: def test_a_volume_comes_back_as_the_dtype_it_went_in_as(dtype: torch.dtype) -> None: """The sampler computes in whatever it must and casts back once. A store's dtype is the store's.""" volume = torch.from_numpy((_volume() % 120).astype(np.float32)).to(dtype) - assert _offset_region(volume).dtype is dtype + assert _sample(volume, mode="linear").dtype is dtype + assert _sample(volume, mode="nearest").dtype is dtype def test_nearest_is_itk_round_half_up_and_not_a_size_ratio() -> None: """``floor(c + 0.5)``, which is a statement about a coordinate. ``F.interpolate``'s nearest is ``floor(o * scale)``, a statement about a size RATIO -- it says - nothing once the target grid carries an origin of its own, which is the whole point of an offset - map. A label map is still a label map under either rule, so only this catches it. + nothing once the target grid carries an origin of its own, and it lags the LINEAR map of the same + stage by ``(scale - 1) / 2`` source voxels, so an image and its label map resampled together come + out shifted against each other. A label map is still a label map under either rule. """ labels = torch.arange(int(np.prod(_SOURCE)), dtype=torch.uint8).reshape(1, *_SOURCE) % 7 - got = _offset_region(labels).numpy()[0] + got = _sample(labels).numpy()[0] expected = np.empty_like(got) source = labels.numpy()[0] @@ -138,11 +129,14 @@ def test_inside_is_the_half_open_half_voxel_rim() -> None: volume = torch.full((1, 4, 4, 4), 5.0) fill = -99.0 - # A target of one voxel per axis, placed by the offset alone. def at(offset: float) -> float: - one = (slice(0, 1), slice(0, 1), slice(0, 1)) - got = _offset_region( - volume, fill=fill, target_slices=one, scales=[1.0, 1.0, 1.0], n_in=[4, 4, 4], offsets=[offset] * 3 + got = _sample( + volume, + fill=fill, + target_shape=(1, 1, 1), + scales=[1.0, 1.0, 1.0], + offsets=[offset] * 3, + source_shape=(4, 4, 4), ) return float(got.flatten()[0]) @@ -152,7 +146,7 @@ def at(offset: float) -> float: assert at(3.5) == fill, "n - 0.5 itself is outside: the interval is half open" -def test_the_separable_sampler_matches_simpleitk() -> None: +def test_the_gather_matches_simpleitk() -> None: """The independent check. Written against SimpleITK because that is what the arithmetic claims. The oracle is skipped here and not at module scope: the rules above -- the working dtype, the @@ -161,15 +155,38 @@ def test_the_separable_sampler_matches_simpleitk() -> None: """ sitk = pytest.importorskip("SimpleITK") volume = _volume() - got = _offset_region(torch.from_numpy(volume)).numpy()[0] + got = _sample(torch.from_numpy(volume)).numpy()[0] image = sitk.GetImageFromArray(volume[0]) image.SetSpacing((1.0, 1.0, 1.0)) image.SetOrigin((0.0, 0.0, 0.0)) - grid = sitk.Image(*reversed([sl.stop - sl.start for sl in _TARGET]), sitk.sitkFloat32) + grid = sitk.Image(*reversed(list(_TARGET)), sitk.sitkFloat32) # sitk takes geometry in (x, y, z) where the arrays above are (z, y, x). grid.SetSpacing(tuple(reversed(_SCALES))) grid.SetOrigin(tuple(reversed(_OFFSETS))) want = sitk.GetArrayFromImage(sitk.Resample(image, grid, sitk.Transform(), sitk.sitkLinear, 0.0)) np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) + + +def test_a_region_reads_the_same_voxels_as_the_whole_volume() -> None: + """Coordinates are GLOBAL, so handing the gather a window changes almost nothing. + + Almost, and the exception is named: a blend goes to ``grid_sample``, which takes NORMALISED + coordinates and therefore divides by the extent of the tensor it is handed -- a window, here. + That is the one region-local number in the path, and it is what a fused kernel costs. Every + other part of the arithmetic is global, which is why the disagreement stays at rounding rather + than moving a sample. + """ + volume = torch.from_numpy(_volume()) + whole = gather(volume, _coordinates(), [0, 0, 0], list(_SOURCE), "linear", 0.0) + + start = [2, 3, 4] + window = volume[:, start[0] :, start[1] :, start[2] :] + partial = gather(window, _coordinates(), start, list(_SOURCE), "linear", 0.0) + + # Only where the whole-volume answer read from inside the window can the two agree; elsewhere the + # window simply does not hold the voxels, which is the caller's contract to respect. + reach = (slice(None), slice(3, None), slice(4, None), slice(5, None)) + span = float(whole.max() - whole.min()) + torch.testing.assert_close(partial[reach], whole[reach], rtol=0, atol=1e-5 * span) diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 969cd7fa..95c05684 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -32,8 +32,10 @@ import torch from konfai.data.case_reduction import CaseReduction from konfai.data.data_manager import _check_patch_transform_locality +from konfai.data.geometry import AffineMap, Grid, TransformBound from konfai.data.patching import DatasetManager, DatasetPatch -from konfai.data.transform import LocalityKind, Reduce, Resample, ResampleToReference, Write +from konfai.data.sampling import source_window +from konfai.data.transform import LocalityKind, Reduce, ResampleToReference, ResampleToShape, Write from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ConfigError, TransformError from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE @@ -149,28 +151,45 @@ def test_the_edge_of_the_data_is_where_simpleitk_puts_it(dataset: Dataset) -> No 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: +def test_a_label_map_takes_the_nearest_voxel_itk_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. + THE ORACLE IS ITK'S EXACT PATH, and it has to be named because ITK ships two. Given a LINEAR + transform, ``ResampleImageFilter`` computes the source index once per scanline and then walks it + by a constant delta; the accumulation drifts, so a continuous index that is exactly ``k + 0.5`` + lands a hair below and rounds down. Given anything it classes non-linear it calls + ``TransformPhysicalPointToContinuousIndex`` per voxel and rounds up, as the rule says. On this + fixture the two disagree with EACH OTHER on 56 of 1050 voxels -- asserted below, so this is a + statement about ITK and not a guess -- and KonfAI matches the exact one. + + Reproducing the drift instead is not available: it depends on where a scanline starts, so a + streamed region and the whole volume would round differently at the same voxel. """ 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)) + fast = sitk.GetArrayFromImage(sitk.Resample(image, grid, sitk.Transform(), sitk.sitkNearestNeighbor, 0)) + # A field of zeros is the identity map that ITK nonetheless classes non-linear, which is the only + # way to ask ResampleImageFilter for its per-voxel arithmetic. + zeros = sitk.GetImageFromArray(np.zeros((*_REFERENCE_SPATIAL, 3), dtype=np.float64), isVector=True) + zeros.SetOrigin(_REFERENCE_ORIGIN) + zeros.SetSpacing(_REFERENCE_SPACING) + exact = sitk.GetArrayFromImage( + sitk.Resample(image, grid, sitk.DisplacementFieldTransform(zeros), sitk.sitkNearestNeighbor, 0) + ) + assert np.count_nonzero(fast != exact) > 0, "the fixture no longer exercises an exact half" assert got.dtype == np.uint8 - np.testing.assert_array_equal(got, want) + np.testing.assert_array_equal(got, exact) @pytest.mark.parametrize("dtype", [np.uint16, np.int16, np.float32]) @@ -330,22 +349,29 @@ def test_a_region_off_the_source_reads_one_voxel(offset: float) -> None: 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)) + source = Grid(_SOURCE_SPATIAL, np.zeros(3), np.ones(3), np.eye(3)) + far = Grid(tuple([4] * len(_SOURCE_SPATIAL)), np.full(3, offset), np.ones(3), np.eye(3)) + window = source_window(far, source, TransformBound.exact(AffineMap.identity(3))) + + assert [part.stop - part.start for part in window] == [1, 1, 1] + assert all(0 <= part.start < extent for part, 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_declares_regrid_for_a_case_that_has_a_geometry(dataset: Dataset) -> None: + """A size ratio would lose the origin entirely and silently, so the stage owns its own map. + + Read off the header in hand: with no geometry there is no physical space to resample in, and + the stage says so instead of declaring a region it could not compute. + """ + stage = _stage(dataset) + assert stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind is LocalityKind.REGRID + + cold = stage.patch_locality(Attribute()) + assert cold.kind is LocalityKind.WHOLE_VOLUME + assert cold.reason is not None and "physical space" in cold.reason def test_it_is_refused_as_a_patch_transform(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: @@ -355,22 +381,36 @@ def test_it_is_refused_as_a_patch_transform(dataset: Dataset, monkeypatch: pytes 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"): + # Config time has no case, so a target that needs a physical space answers WHOLE_VOLUME -- still + # refused, and by the row that says so. + with pytest.raises(ConfigError, match="needs the whole volume"): _check_patch_transform_locality(_stage(dataset), "CT", "CT") + # A change of extent needs no geometry at all, so it reaches config time as what it is. + with pytest.raises(ConfigError, match="onto another grid"): + _check_patch_transform_locality(ResampleToShape(shape=[4, 4, 4]), "CT", "CT") + + +def test_a_differing_direction_is_resampled_and_not_refused(tmp_path: Path) -> None: + """Axes that do not line up used to be refused. They are a rotation, and a rotation is ordinary. -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.""" + The old map was a scale and a shift per axis, which no rotation is, so the stage refused and told + the reader to run ``Canonical`` first -- a second resample, and a second interpolation of the + same voxels. The source region of a target region is now that region's world box mapped through + the map and read back as an index window, which a rotation answers as readily as a translation. + """ 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) - ) + volume = _volume(_SOURCE_SPATIAL) + reference = _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING, turned) + dataset.write("Case", _CASE, volume, source) + dataset.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), reference) - with pytest.raises(TransformError, match="Direction cosines differ"): - _stage(dataset).transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), source) + stage = _stage(dataset) + assert stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), Attribute(source)) == list(_REFERENCE_SPATIAL) + got = stage(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy()[0] + np.testing.assert_allclose(got, _simpleitk(volume, source, reference), rtol=1e-5, atol=1e-3) def test_a_case_with_no_geometry_is_refused(dataset: Dataset) -> None: @@ -401,8 +441,8 @@ def test_a_case_that_never_meets_the_reference_is_refused(tmp_path: Path) -> Non 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() + with pytest.raises(TransformError, match="cannot find reference 'NOT_THERE'"): + stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) def test_an_unnamed_group_is_refused_when_the_store_has_several(dataset: Dataset) -> None: @@ -410,7 +450,7 @@ def test_an_unnamed_group_is_refused_when_the_store_has_several(dataset: Dataset stage = ResampleToReference(entry=_CASE) stage.set_datasets([dataset]) with pytest.raises(TransformError, match="cannot tell which group"): - stage.reference_grid() + stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) def test_an_empty_entry_is_refused_at_construction() -> None: @@ -471,11 +511,32 @@ def test_the_inverse_returns_the_case_to_its_own_grid(dataset: Dataset) -> None: 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.""" +def test_the_inverse_streams_once_the_forward_has_stacked_its_geometry(dataset: Dataset) -> None: + """A change of grid inverts to a change of grid, and needs no memory of the case to do it. + + The forward stacks the target geometry over the source's, so the inverse reads the grid it is + holding, pops, and reads the grid it is restoring -- both off the attribute in hand. That is what + lets a prediction finalize through this stage stay streamed instead of assembling the volume. + """ + stage = _stage(dataset) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + stage(_CASE, torch.from_numpy(_volume(_SOURCE_SPATIAL).copy()), source) + + locality = stage.inverse_patch_locality(Attribute(source)) + assert locality.kind is LocalityKind.REGRID + assert stage.inverse_transform_shape(list(_REFERENCE_SPATIAL), Attribute(source)) == list(_SOURCE_SPATIAL) + + restored = stage.inverse(_CASE, torch.zeros(1, *_REFERENCE_SPATIAL), source) + assert list(restored.shape[1:]) == list(_SOURCE_SPATIAL) + np.testing.assert_allclose(source.get_np_array("Spacing"), _SOURCE_SPACING) + np.testing.assert_allclose(source.get_np_array("Origin"), _SOURCE_ORIGIN) + + +def test_the_inverse_says_why_when_the_forward_left_no_stack(dataset: Dataset) -> None: + """Asked cold, it cannot know the shape it restores -- and says that rather than guessing.""" 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 + assert locality.reason is not None and "not on the attribute" in locality.reason # --------------------------------------------------------------------- through a field @@ -538,6 +599,7 @@ def _simpleitk_warp( interpolator: int = sitk.sitkLinear, pixel: int = sitk.sitkFloat32, fill: float = _FILL, + field_direction: np.ndarray | None = None, ) -> np.ndarray: """``sitk.Resample(image, grid, DisplacementFieldTransform(field))`` — the one-pass authority.""" image = sitk.GetImageFromArray(volume[0]) @@ -552,6 +614,8 @@ def _simpleitk_warp( vector = sitk.GetImageFromArray(np.moveaxis(field, 0, -1).astype(np.float64), isVector=True) vector.SetOrigin(_FIELD_ORIGIN) vector.SetSpacing(_FIELD_SPACING) + if field_direction is not None: + vector.SetDirection(field_direction.reshape(-1).tolist()) transform = sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)) return sitk.GetArrayFromImage(sitk.Resample(image, grid, transform, interpolator, fill)) @@ -821,15 +885,28 @@ def _stage_regrid_kind(images: Dataset) -> LocalityKind: 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 +def test_a_field_on_another_direction_is_read_where_it_is_asked( + warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path +) -> None: + """A field stored on rotated axes is still a displacement in world units, and is read as one. + + It used to be refused for the same reason a rotated reference was: the map onto the field's grid + was a per-axis scale and shift. It is now the field's own ``world -> index``, so where the field + was stored stops being a constraint on where it can be applied. + """ + 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)) + field = _displacement() + turned.write("DVF", _CASE, field, _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)) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + + assert stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), Attribute(source)) == list(_REFERENCE_SPATIAL) + got = stage(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy()[0] + + want = _simpleitk_warp(volume, field, field_direction=rotated) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-3) def test_fields_can_live_beside_the_cases(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: @@ -881,16 +958,22 @@ def test_a_bound_with_no_field_is_refused() -> None: ResampleToReference(entry=_CASE, group="Reference", max_displacement=1.0) -def test_the_two_gathers_agree_bit_for_bit_through_an_identity_field(tmp_path: Path) -> None: +def test_the_two_gathers_obey_the_same_rules_through_an_identity_field(tmp_path: Path) -> None: """One arithmetic, two loops: per-axis maps, and eight corners at a coordinate volume. - The separable loop cannot serve a displacement (a displacement is not separable) and the flat - gather is the slower way to do a map that is. So both exist, and both have to obey the same - inside interval, the same tap clamp and the same fill. A ZERO field is where that is checkable: - the composed path reduces to the grid change alone, so the two must land on the same voxels. - - Bit for bit, not close. A tolerance here would hide exactly the drift this guards -- one loop - keeping a rule the other quietly dropped, which is how the CPU-half guard went missing once. + The separable loop cannot serve a displacement (a displacement is not separable) and the corner + gather is far the slower way to do a map that is -- measured at 43x versus 660x of + ``F.interpolate`` on a CT. So both exist, and both have to obey the same inside interval, the + same tap clamp and the same fill. A ZERO field is where that is checkable: the composed path + reduces to the grid change alone, so the two must land on the same voxels. + + WHICH VOXELS ARE FILL is asserted exactly, because that is the rule a drifting loop drops -- one + of them quietly widening its domain by half a voxel is invisible in the values and obvious here. + The VALUES agree to float rounding and not bit for bit: the separable loop sums the tensor + product axis by axis and the other corner by corner, which is the same sum in a different order. + Nothing needs them identical -- a map either factorises or it does not, so no case is ever served + by both -- and the equality that must be exact, a streamed region against the whole volume, is + asserted where it lives. """ images = Dataset(tmp_path / "Images", "h5") volume = _high_frequency() @@ -904,4 +987,6 @@ def test_the_two_gathers_agree_bit_for_bit_through_an_identity_field(tmp_path: P composed = _warping(images, fields, fill=_FILL)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy() assert 0 < int((separable == _FILL).sum()) < separable.size, "the fixture must have a rim, and not be all rim" - np.testing.assert_array_equal(composed, separable) + np.testing.assert_array_equal(composed == _FILL, separable == _FILL) + span = float(separable.max() - separable.min()) + np.testing.assert_allclose(composed, separable, rtol=0, atol=1e-5 * span) diff --git a/tests/unit/test_resample_transform.py b/tests/unit/test_resample_transform.py new file mode 100644 index 00000000..67b3056c --- /dev/null +++ b/tests/unit/test_resample_transform.py @@ -0,0 +1,340 @@ +# 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 + +"""``ResampleTransform`` streaming: against SimpleITK, and against its own whole-volume path. + +The fixture is high-frequency and its direction cosines are oblique, because a smooth phantom on +an axis-aligned grid passes a map that is wrong in exactly the ways this stage can be wrong. +""" + +import numpy as np +import pytest +import torch +from konfai.data.transform import LocalityKind, RegionContext, ResampleTransform +from konfai.utils.dataset import Attribute +from konfai.utils.errors import TransformError + +sitk = pytest.importorskip("SimpleITK") + +SIZE = (22, 28, 34) +CASE = "CASE_000" + + +def _phantom() -> np.ndarray: + z, y, x = np.meshgrid(*[np.arange(extent, dtype=np.float64) for extent in SIZE], indexing="ij") + return (100.0 * np.sin(1.7 * z) * np.cos(2.1 * y) + 80.0 * np.sin(2.9 * x)).astype(np.float32) + + +def _image(oblique: bool = True) -> "sitk.Image": + image = sitk.GetImageFromArray(_phantom()) + image.SetSpacing((0.8, 1.2, 1.5)) + image.SetOrigin((10.0, -5.0, 2.0)) + if oblique: + a, b = np.deg2rad(20.0), np.deg2rad(15.0) + rz = np.array([[np.cos(a), -np.sin(a), 0.0], [np.sin(a), np.cos(a), 0.0], [0.0, 0.0, 1.0]]) + ry = np.array([[np.cos(b), 0.0, np.sin(b)], [0.0, 1.0, 0.0], [-np.sin(b), 0.0, np.cos(b)]]) + image.SetDirection(tuple((rz @ ry).ravel())) + return image + + +def _attribute(image: "sitk.Image") -> Attribute: + attribute = Attribute() + attribute["Origin"] = np.asarray(image.GetOrigin()) + attribute["Spacing"] = np.asarray(image.GetSpacing()) + attribute["Direction"] = np.asarray(image.GetDirection()) + return attribute + + +class _StoredTransform: + """The smallest thing that answers what ``ResampleTransform`` asks of a ``Dataset``.""" + + def __init__(self, group: str, transform: "sitk.Transform") -> None: + self.group = group + self.transform = transform + + def is_dataset_exist(self, group: str, name: str) -> bool: + del name + return group == self.group + + def read_transform(self, group: str, name: str) -> "sitk.Transform": + del group, name + return self.transform + + +def _euler(image): + transform = sitk.Euler3DTransform() + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetRotation(0.11, -0.2, 0.31) + transform.SetTranslation((3.0, -2.0, 1.0)) + return transform + + +def _affine(image): + transform = sitk.AffineTransform(3) + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetMatrix(np.array([[1.1, 0.05, 0.0], [0.0, 0.9, 0.07], [0.02, 0.0, 1.2]]).ravel()) + transform.SetTranslation((4.0, -3.0, 2.0)) + return transform + + +def _bspline(image, amplitude: float = 7.0): + transform = sitk.BSplineTransformInitializer(image, [5] * 3, 3) + size = np.asarray(transform.GetParameters()).size + transform.SetParameters(list(np.random.RandomState(2).uniform(-amplitude, amplitude, size))) + return transform + + +def _field(image): + filt = sitk.TransformToDisplacementFieldFilter() + filt.SetReferenceImage(image) + return sitk.DisplacementFieldTransform(sitk.Cast(filt.Execute(_bspline(image)), sitk.sitkVectorFloat64)) + + +def _families(image) -> list[tuple[str, "sitk.Transform"]]: + return [ + ("euler", _euler(image)), + ("affine", _affine(image)), + ("bspline", _bspline(image)), + ("field", _field(image)), + ] + + +def _stage(image, transform, **kwargs) -> ResampleTransform: + stage = ResampleTransform(transforms={"reg": False}, **kwargs) + stage.set_datasets([_StoredTransform("reg", transform)]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + return stage + + +DEVICES = [torch.device("cpu")] + ([torch.device("cuda")] if torch.cuda.is_available() else []) +DEVICE_IDS = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +@pytest.mark.parametrize("oblique", [False, True], ids=["axis-aligned", "oblique"]) +def test_the_whole_volume_path_matches_simpleitk(oblique: bool): + device = torch.device("cpu") + image = _image(oblique) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0).to(device) + for label, transform in _families(image): + stage = _stage(image, transform) + want = sitk.GetArrayFromImage(sitk.Resample(image, image, transform, sitk.sitkLinear, 0.0)) + got = stage(CASE, volume, _attribute(image)).squeeze(0).cpu().numpy() + deviation = float(np.abs(want - got).max()) + assert deviation <= 1e-3 * float(np.abs(want).max()), f"{label}: {deviation:.4g}" + + +@pytest.mark.parametrize("device", DEVICES, ids=DEVICE_IDS) +@pytest.mark.parametrize("rows", [3, 8], ids=["slab-3", "slab-8"]) +def test_the_streamed_slabs_agree_with_the_whole_volume(device: torch.device, rows: int): + """A slab puts every sample where the whole volume puts it, to a fraction of the data's range. + + A stored transform never factorises, so this is the general path: the blend goes to + ``grid_sample``, one fused kernel worth 4x, which takes NORMALISED coordinates and so divides by + the extent of the tensor handed to it -- and a slab is handed a window. That single region-local + number is the whole of the disagreement. Everything else in the path is global, which is why it + stays at rounding instead of moving a sample. + + The bound is far above what is measured and far below a moved map, which is wrong by voxels. + Bit-identity still holds on the separable path, and ``test_resample.py`` pins it there. + """ + image = _image() + attribute = _attribute(image) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0).to(device) + for label, transform in _families(image): + stage = _stage(image, transform) + reference = stage(CASE, volume, Attribute(attribute)) + streamed = torch.empty_like(reference) + for start in range(0, SIZE[0], rows): + stop = min(start + rows, SIZE[0]) + target = (slice(start, stop), slice(0, SIZE[1]), slice(0, SIZE[2])) + source = tuple(stage.stream_region_source(CASE, target, list(SIZE), Attribute(attribute))) + block = volume[(slice(None), *source)] + context = RegionContext(source, target, tuple(SIZE), tuple(SIZE)) + streamed[(slice(None), *target)] = stage.stream_region(CASE, block, context, Attribute(attribute)) + span = float((reference.max() - reference.min()).item()) + torch.testing.assert_close(streamed, reference, rtol=0.0, atol=1e-5 * span, msg=label) + + +def _slab_reads(stage: ResampleTransform, attribute: Attribute, rows: int) -> list[int]: + reads = [] + for start in range(0, SIZE[0], rows): + target = (slice(start, min(start + rows, SIZE[0])), slice(0, SIZE[1]), slice(0, SIZE[2])) + source = stage.stream_region_source(CASE, target, list(SIZE), Attribute(attribute)) + reads.append(int(np.prod([part.stop - part.start for part in source]))) + return reads + + +def test_a_slab_pulls_a_bounded_source_window(): + # The point of the whole exercise: on a map aligned with the storage axes a slab reads a slab, + # not the volume. + image = _image(oblique=False) + stage = _stage(image, sitk.TranslationTransform(3, (2.0, -1.5, 1.0))) + attribute = _attribute(image) + for start in range(0, SIZE[0], 4): + target = (slice(start, min(start + 4, SIZE[0])), slice(0, SIZE[1]), slice(0, SIZE[2])) + source = stage.stream_region_source(CASE, target, list(SIZE), Attribute(attribute)) + assert source[0].stop - source[0].start <= 8, "a translated slab pulled more than its neighbours" + assert sum(_slab_reads(stage, attribute, 4)) < 3 * np.prod(SIZE) + + +def test_an_oblique_map_does_not_bound_and_the_numbers_say_so(): + # Not a defect of the bound -- the bound is exact here. A thin slab rotated against the storage + # axes genuinely has an axis-aligned source box covering most of the volume, and it gets worse + # the finer the decomposition. This is the measurement a cost model exists to report, and the + # reason streaming this map is not automatically worth it. + image = _image(oblique=True) + stage = _stage(image, _euler(image)) + attribute = _attribute(image) + coarse = sum(_slab_reads(stage, attribute, SIZE[0])) / np.prod(SIZE) + fine = sum(_slab_reads(stage, attribute, 4)) / np.prod(SIZE) + assert fine > coarse, f"amplification must grow as slabs thin: {coarse:.2f}x then {fine:.2f}x" + assert fine > 2.0, f"the oblique fixture is meant to be expensive, got {fine:.2f}x" + + +class TestLocality: + def test_a_boundable_cohort_declares_regrid(self): + image = _image() + for label, transform in _families(image): + stage = _stage(image, transform) + assert stage.patch_locality(_attribute(image)).kind is LocalityKind.REGRID, label + + def test_a_case_without_geometry_falls_back_and_says_why(self): + image = _image() + stage = ResampleTransform(transforms={"reg": False}) + stage.set_datasets([_StoredTransform("reg", _euler(image))]) + stage.transform_shape("", CASE, list(SIZE), Attribute()) # no Origin/Spacing/Direction + locality = stage.patch_locality(Attribute()) # judged on the header handed over + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert "physical space" in (locality.reason or "") + + def test_a_missing_transform_falls_back_and_says_which_group(self): + image = _image() + stage = ResampleTransform(transforms={"absent": False}) + stage.set_datasets([_StoredTransform("reg", _euler(image))]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + locality = stage.patch_locality(_attribute(image)) + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert "absent" in (locality.reason or "") + + def test_a_spline_order_with_no_kernel_falls_back_instead_of_crashing_mid_run(self): + """ITK writes orders 0 and 2 as readily as 3, and neither has a kernel here. + + The refusal has to happen where the value is BUILT, not where it is finally sampled: a stage + that decodes such a spline without complaint declares REGRID, passes the plan, and raises on + the first region -- which is halfway through a run, per case, after bytes are already + written. Refused at decode, it is one more whole-volume line in the plan. + """ + image = _image() + quadratic = sitk.BSplineTransformInitializer(image, [5] * 3, 2) + size = np.asarray(quadratic.GetParameters()).size + quadratic.SetParameters(list(np.random.RandomState(3).uniform(-5.0, 5.0, size))) + + stage = ResampleTransform(transforms={"reg": False}) + stage.set_datasets([_StoredTransform("reg", quadratic)]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + + locality = stage.patch_locality(_attribute(image)) + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert "order 2" in (locality.reason or "") + + def test_inverting_a_spline_falls_back_with_the_remedy(self): + image = _image() + stage = ResampleTransform(transforms={"reg": True}) + stage.set_datasets([_StoredTransform("reg", _bspline(image))]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + locality = stage.patch_locality(_attribute(image)) + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert "Store the inverse" in (locality.reason or "") + + def test_inverting_a_rigid_map_is_exact_and_still_streams(self): + image = _image() + transform = _euler(image) + stage = ResampleTransform(transforms={"reg": True}) + stage.set_datasets([_StoredTransform("reg", transform)]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + assert stage.patch_locality(_attribute(image)).kind is LocalityKind.REGRID + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0) + want = sitk.GetArrayFromImage(sitk.Resample(image, image, transform.GetInverse(), sitk.sitkLinear, 0.0)) + got = stage(CASE, volume, _attribute(image)).squeeze(0).numpy() + assert float(np.abs(want - got).max()) <= 1e-3 * float(np.abs(want).max()) + + +class TestSampling: + def test_a_label_map_is_not_blended(self): + labels = (np.abs(_phantom()) % 5).astype(np.uint8) + image = sitk.GetImageFromArray(labels) + image.SetSpacing((0.8, 1.2, 1.5)) + image.SetOrigin((10.0, -5.0, 2.0)) + stage = _stage(image, _euler(image)) + got = stage(CASE, torch.from_numpy(labels).unsqueeze(0), _attribute(image)).squeeze(0).numpy() + want = sitk.GetArrayFromImage(sitk.Resample(image, image, _euler(image), sitk.sitkNearestNeighbor, 0.0)) + np.testing.assert_array_equal(got, want) + assert set(np.unique(got)) <= set(np.unique(labels)) + + def test_the_fill_reaches_where_the_map_leaves_the_source(self): + image = _image(oblique=False) + transform = sitk.TranslationTransform(3, (12.0, 9.0, 6.0)) + stage = _stage(image, transform, fill=-1234.0) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0) + got = stage(CASE, volume, _attribute(image)).squeeze(0).numpy() + want = sitk.GetArrayFromImage(sitk.Resample(image, image, transform, sitk.sitkLinear, -1234.0)) + assert 0 < int((want == -1234.0).sum()) < want.size + np.testing.assert_array_equal(got == -1234.0, want == -1234.0) + + +class TestRefusals: + def test_an_unknown_interpolation_is_refused_at_construction(self): + with pytest.raises(TransformError, match="interpolation"): + ResampleTransform(transforms={"reg": False}, interpolation="bspline") + + def test_no_transforms_is_refused_at_construction(self): + with pytest.raises(TransformError, match="at least one group"): + ResampleTransform(transforms={}) + + def test_the_inverse_direction_says_what_to_do_instead(self): + stage = ResampleTransform(transforms={"reg": False}) + assert stage.inverse_patch_locality(Attribute()).kind is LocalityKind.WHOLE_VOLUME + with pytest.raises(TransformError, match="inverse: false"): + stage.inverse(CASE, torch.zeros(1, 2, 2, 2), Attribute()) + + +class TestCompositeOrder: + def test_two_groups_compose_as_they_always_did(self): + # Declaration order has always meant SimpleITK's composite order (last declared applied + # first), because this stage built a CompositeTransform from the declared list. Decoding + # normalizes to application order, so the reversal has to be reinstated -- and pinned. + image = _image(oblique=False) + first = sitk.TranslationTransform(3, (4.0, 0.0, 0.0)) + second = sitk.ScaleTransform(3, (1.5, 1.5, 1.5)) + stage = ResampleTransform(transforms={"a": False, "b": False}) + + class _Two: + def is_dataset_exist(self, group: str, name: str) -> bool: + del name + return group in ("a", "b") + + def read_transform(self, group: str, name: str): + del name + return first if group == "a" else second + + stage.set_datasets([_Two()]) + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0) + got = stage(CASE, volume, _attribute(image)).squeeze(0).numpy() + want = sitk.GetArrayFromImage( + sitk.Resample(image, image, sitk.CompositeTransform([first, second]), sitk.sitkLinear, 0.0) + ) + assert float(np.abs(want - got).max()) <= 1e-3 * float(np.abs(want).max()) diff --git a/tests/unit/test_sampling.py b/tests/unit/test_sampling.py new file mode 100644 index 00000000..8b34fc05 --- /dev/null +++ b/tests/unit/test_sampling.py @@ -0,0 +1,280 @@ +# 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 + +"""The torch coordinate producer and gather, against SimpleITK. + +The oracle is external on purpose: two KonfAI paths agree by construction, including on a grid +placed in the wrong place. The fixture is high-frequency — a smooth phantom would pass a wrong +map — and its direction cosines are oblique, which is the case the separable samplers refuse and +this one exists for. +""" + +import numpy as np +import pytest +import torch +from konfai.data.geometry import Grid, bound_of +from konfai.data.sampling import gather, read_amplification, source_index, source_window + +sitk = pytest.importorskip("SimpleITK") + +from konfai.utils.ITK import decode_transform_stages # noqa: E402 + +SIZE = (24, 30, 36) + + +def _phantom(size=SIZE) -> np.ndarray: + z, y, x = np.meshgrid(*[np.arange(extent, dtype=np.float64) for extent in size], indexing="ij") + return (100.0 * np.sin(1.7 * z) * np.cos(2.1 * y) + 80.0 * np.sin(2.9 * x)).astype(np.float32) + + +def _image(oblique: bool = True, size=SIZE) -> "sitk.Image": + image = sitk.GetImageFromArray(_phantom(size)) + image.SetSpacing((0.8, 1.2, 1.5)) + image.SetOrigin((10.0, -5.0, 2.0)) + if oblique: + a, b = np.deg2rad(20.0), np.deg2rad(15.0) + rz = np.array([[np.cos(a), -np.sin(a), 0.0], [np.sin(a), np.cos(a), 0.0], [0.0, 0.0, 1.0]]) + ry = np.array([[np.cos(b), 0.0, np.sin(b)], [0.0, 1.0, 0.0], [-np.sin(b), 0.0, np.cos(b)]]) + image.SetDirection(tuple((rz @ ry).ravel())) + return image + + +def _grid(image: "sitk.Image") -> Grid: + rank = image.GetDimension() + return Grid( + tuple(int(extent) for extent in reversed(image.GetSize())), + np.asarray(image.GetOrigin(), dtype=np.float64), + np.asarray(image.GetSpacing(), dtype=np.float64), + np.asarray(image.GetDirection(), dtype=np.float64).reshape(rank, rank), + ) + + +def _euler(image): + transform = sitk.Euler3DTransform() + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetRotation(0.11, -0.2, 0.31) + transform.SetTranslation((3.0, -2.0, 1.0)) + return transform + + +def _affine(image): + transform = sitk.AffineTransform(3) + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetMatrix(np.array([[1.1, 0.05, 0.0], [0.0, 0.9, 0.07], [0.02, 0.0, 1.2]]).ravel()) + transform.SetTranslation((4.0, -3.0, 2.0)) + return transform + + +def _bspline(image, mesh: int = 5, amplitude: float = 7.0): + transform = sitk.BSplineTransformInitializer(image, [mesh] * 3, 3) + size = np.asarray(transform.GetParameters()).size + transform.SetParameters(list(np.random.RandomState(2).uniform(-amplitude, amplitude, size))) + return transform + + +def _field(image): + filt = sitk.TransformToDisplacementFieldFilter() + filt.SetReferenceImage(image) + return sitk.DisplacementFieldTransform(sitk.Cast(filt.Execute(_bspline(image)), sitk.sitkVectorFloat64)) + + +def _transforms(image) -> list[tuple[str, "sitk.Transform"]]: + return [ + ("euler", _euler(image)), + ("affine", _affine(image)), + ("bspline", _bspline(image)), + ("field", _field(image)), + ("composite", sitk.CompositeTransform([_affine(image), _bspline(image)])), + ] + + +def _resample_whole(image, transform, interpolator=sitk.sitkLinear, fill: float = 0.0) -> np.ndarray: + return sitk.GetArrayFromImage(sitk.Resample(image, image, transform, interpolator, fill)) + + +def _konfai_region( + image, transform, region: tuple[slice, ...], device: torch.device, mode: str = "linear", fill: float = 0.0 +) -> np.ndarray: + """One target region through the torch path, reading only the bounded source window.""" + grid = _grid(image) + stages = decode_transform_stages(transform) + bound = bound_of(stages, 3) + target = grid.sub_grid(region) + window = source_window(target, grid, bound) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0) + block = volume[(slice(None), *window)].to(device) + coordinates = source_index(target, grid, stages, device) + out = gather( + block, + coordinates, + [part.start for part in window], + list(grid.size_zyx), + mode, + fill, + ) + return out.squeeze(0).cpu().numpy() + + +DEVICES = [torch.device("cpu")] + ([torch.device("cuda")] if torch.cuda.is_available() else []) +DEVICE_IDS = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +@pytest.mark.parametrize("oblique", [False, True], ids=["axis-aligned", "oblique"]) +def test_the_whole_grid_matches_simpleitk(oblique: bool): + device = torch.device("cpu") + image = _image(oblique) + whole = (slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2])) + for label, transform in _transforms(image): + want = _resample_whole(image, transform) + got = _konfai_region(image, transform, whole, device) + deviation = float(np.abs(want - got).max()) + scale = float(np.abs(want).max()) + assert deviation <= 1e-3 * scale, f"{label}: {deviation:.4g} against a range of {scale:.4g}" + + +@pytest.mark.parametrize("device", DEVICES, ids=DEVICE_IDS) +@pytest.mark.parametrize("rows", [3, 7, SIZE[0]], ids=["slab-3", "slab-7", "whole"]) +def test_the_streamed_slabs_agree_with_the_whole_volume(device: torch.device, rows: int): + """A slab and the whole volume put every sample in the same place, to a stated bound. + + NOT bit for bit, and the reason is deliberate: a blend through a map that does not factorise + goes to ``grid_sample``, one fused kernel worth 4x on a warp, and grid_sample takes NORMALISED + coordinates -- so it divides by the extent of the tensor handed to it, and a slab is handed a + window. What is bit-identical is the SEPARABLE path, which is most resamples, and which + ``test_resample.py`` pins. + + The bound below is a fraction of the data's own range, and it is roughly thirty times the worst + disagreement measured -- while a slab whose map actually moved is wrong by VOXELS, orders of + magnitude above it. The companion test underneath keeps that end honest. + """ + image = _image() + whole_region = (slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2])) + for label, transform in _transforms(image): + reference = _konfai_region(image, transform, whole_region, device) + streamed = np.empty_like(reference) + for start in range(0, SIZE[0], rows): + stop = min(start + rows, SIZE[0]) + region = (slice(start, stop), slice(0, SIZE[1]), slice(0, SIZE[2])) + streamed[start:stop] = _konfai_region(image, transform, region, device) + span = float(reference.max() - reference.min()) + np.testing.assert_allclose(streamed, reference, rtol=0, atol=1e-5 * span, err_msg=label) + + +def test_a_slab_left_at_the_volume_origin_is_loudly_wrong(): + device = torch.device("cpu") + # The silent failure mode, tested first: a region whose grid keeps the VOLUME's origin replays + # the same part of the source for every slab, and the output still looks like an image. If this + # test ever passes quietly, sub_grid stopped placing regions. + image = _image() + transform = _euler(image) + grid = _grid(image) + stages = decode_transform_stages(transform) + reference = _konfai_region(image, transform, (slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2])), device) + + rows = 6 + wrong = np.empty_like(reference) + volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0).to(device) + for start in range(0, SIZE[0], rows): + stop = min(start + rows, SIZE[0]) + misplaced = Grid((stop - start, SIZE[1], SIZE[2]), grid.origin_xyz, grid.spacing_xyz, grid.direction_xyz) + coordinates = source_index(misplaced, grid, stages, device) + wrong[start:stop] = ( + gather(volume, coordinates, [0, 0, 0], list(grid.size_zyx), "linear", 0.0).squeeze(0).cpu().numpy() + ) + assert np.abs(wrong - reference).max() > 0.1 * float(np.abs(reference).max()) + + +def test_the_fill_mask_matches_simpleitk_voxel_for_voxel(): + device = torch.device("cpu") + # The sharp test: where the map leaves the source, SimpleITK writes the fill and so must this. + # A source window one voxel short does not raise — it returns background — so the only thing + # that catches it is comparing WHICH voxels are fill, with no tolerance at all. + image = _image() + transform = sitk.TranslationTransform(3, (14.0, 9.0, 7.0)) + fill = -1234.0 + want = _resample_whole(image, transform, fill=fill) + got = _konfai_region(image, transform, (slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2])), device, fill=fill) + assert 0 < int((want == fill).sum()) < want.size, "the fixture must actually leave the source" + np.testing.assert_array_equal(got == fill, want == fill) + + +def test_nearest_is_byte_identical_on_a_label_map(): + device = torch.device("cpu") + labels = (np.abs(_phantom()) % 7).astype(np.uint8) + image = sitk.GetImageFromArray(labels) + image.SetSpacing((0.8, 1.2, 1.5)) + image.SetOrigin((10.0, -5.0, 2.0)) + transform = _euler(image) + want = _resample_whole(image, transform, interpolator=sitk.sitkNearestNeighbor) + got = _konfai_region( + image, transform, (slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2])), device, mode="nearest" + ) + np.testing.assert_array_equal(got, want) + + +class TestSourceWindow: + def test_the_window_covers_every_coordinate_the_region_samples(self): + image = _image() + grid = _grid(image) + for label, transform in _transforms(image): + stages = decode_transform_stages(transform) + bound = bound_of(stages, 3) + region = (slice(6, 13), slice(0, SIZE[1]), slice(0, SIZE[2])) + target = grid.sub_grid(region) + window = source_window(target, grid, bound) + coordinates = source_index(target, grid, stages, torch.device("cpu")).numpy() + for axis in range(3): + array_axis = 2 - axis + sampled = coordinates[..., axis] + # Only what actually lands on the source has to be covered; the rest takes the fill. + inside = (sampled >= -0.5) & (sampled < grid.size_zyx[array_axis] - 0.5) + if not inside.any(): + continue + low, high = float(sampled[inside].min()), float(sampled[inside].max()) + extent = int(grid.size_zyx[array_axis]) + # The taps a sample needs, clamped to the source exactly as the gather clamps them: + # a coordinate in the half-voxel rim reads the border voxel, it does not read past it. + need_low = min(max(int(np.floor(low)), 0), extent - 1) + need_high = min(max(int(np.floor(high)) + 1, 0), extent - 1) + part = window[array_axis] + assert part.start <= need_low and need_high < part.stop, ( + f"{label}: axis {array_axis} needs [{need_low}, {need_high}] outside {part}" + ) + + +class TestReadAmplification: + def test_it_grows_as_the_decomposition_gets_finer(self): + # The property the cost model rests on, and the reason a gate is needed at all: streaming + # finer never reads less. Measured on an oblique map, where it runs away fastest. + image = _image() + grid = _grid(image) + bound = bound_of(decode_transform_stages(_affine(image)), 3) + ratios = [] + for rows in (SIZE[0], 8, 3, 1): + regions = [ + (slice(start, min(start + rows, SIZE[0])), slice(0, SIZE[1]), slice(0, SIZE[2])) + for start in range(0, SIZE[0], rows) + ] + ratios.append(read_amplification(grid, grid, bound, regions)) + assert ratios == sorted(ratios), f"amplification must be monotone in fineness, got {ratios}" + assert ratios[-1] > 2.0 * ratios[0] + + def test_a_pure_translation_of_a_whole_volume_reads_about_once(self): + image = _image(oblique=False) + grid = _grid(image) + bound = bound_of(decode_transform_stages(sitk.TranslationTransform(3, (1.0, 1.0, 1.0))), 3) + whole = [(slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2]))] + assert read_amplification(grid, grid, bound, whole) == pytest.approx(1.0, abs=0.35) diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 11609e67..4960ba4a 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -45,6 +45,7 @@ Normalize, PatchLocality, Permute, + RegionContext, ResampleToShape, Softmax, TensorCast, @@ -163,7 +164,7 @@ def test_stream_composed_rescale_and_orientation_matches_whole_volume(assert_str volume, [ResampleToShape(shape=[12, 12]), Flip("0")], [4, 4], atol=1e-3 ) plans = manager._resolve_patch_stream_source(0, True).stage_plans - assert [plan.kind.value for plan in plans] == ["rescale", "orientation"] + assert [plan.kind.value for plan in plans] == ["regrid", "orientation"] assert tuple(plans[1].in_shape) == (12, 12) @@ -175,7 +176,7 @@ def test_stream_composed_triple_region_chain_matches_whole_volume(assert_stream_ volume, [Flip("0"), ResampleToShape(shape=[12, 9]), Permute("1|0")], [4, 4], atol=1e-3 ) plans = manager._resolve_patch_stream_source(0, True).stage_plans - assert [plan.kind.value for plan in plans] == ["orientation", "rescale", "orientation"] + assert [plan.kind.value for plan in plans] == ["orientation", "regrid", "orientation"] assert tuple(plans[2].out_shape) == (9, 12) @@ -312,14 +313,15 @@ def test_streamed_nearest_resample_matches_whole_volume_at_any_ratio(n_in: int, expected = resample("case", volume.clone(), Attribute(attribute)) target = tuple(slice(0, n_out) for _ in range(3)) - slices, starts, scales, n_in_list, _ = resample.resample_source_region(target, [n_in] * 3, Attribute(attribute)) - got = resample.resample_region(volume[(slice(None), *slices)], target, starts, scales, n_in_list) + window = resample.stream_region_source("case", target, [n_in] * 3, Attribute(attribute)) + context = RegionContext(tuple(window), target, (n_in,) * 3, (n_out,) * 3) + got = resample.stream_region("case", volume[(slice(None), *window)], context, Attribute(attribute)) torch.testing.assert_close(got, expected, rtol=0, atol=0) @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) def test_streamed_resample_handles_2d(dtype: torch.dtype) -> None: - """resample_region must not assume three spatial axes.""" + """The gather must not assume three spatial axes.""" volume = (torch.arange(1 * 9 * 11, dtype=torch.float32).reshape(1, 9, 11) % 17).to(dtype) resample = ResampleToShape(shape=[5, 6], inverse=False) attribute = Attribute() @@ -327,9 +329,10 @@ def test_streamed_resample_handles_2d(dtype: torch.dtype) -> None: expected = resample("case", volume.clone(), Attribute(attribute)) target = (slice(0, 5), slice(0, 6)) - slices, starts, scales, n_in_list, _ = resample.resample_source_region(target, [9, 11], Attribute(attribute)) - got = resample.resample_region(volume[(slice(None), *slices)], target, starts, scales, n_in_list) - torch.testing.assert_close(got, expected, rtol=0, atol=1e-5) + window = resample.stream_region_source("case", target, [9, 11], Attribute(attribute)) + context = RegionContext(tuple(window), target, (9, 11), (5, 6)) + got = resample.stream_region("case", volume[(slice(None), *window)], context, Attribute(attribute)) + torch.testing.assert_close(got, expected, rtol=0, atol=0) # -------------------------------------------------------------------------------------- @@ -398,7 +401,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: return PatchLocality(LocalityKind.ORIENTATION) def stream_region_source( - self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + self, name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute ) -> list[slice]: return [ slice(extent - t.stop, extent - t.start) diff --git a/tests/unit/test_streamed_write_dispatcher.py b/tests/unit/test_streamed_write_dispatcher.py index 746b279a..9538d83c 100644 --- a/tests/unit/test_streamed_write_dispatcher.py +++ b/tests/unit/test_streamed_write_dispatcher.py @@ -41,6 +41,7 @@ Normalize, Padding, Permute, + RegionContext, ResampleToResolution, Softmax, Standardize, @@ -93,7 +94,7 @@ def test_stream_orientation_mirrored_slab_axis_matches_whole_volume(seed: int) - flip = Flip("0") reference = flip.inverse("case", volume, Attribute()) got = _run_stream( - lambda target: flip.stream_region_target(target, [Z, Y, X], Attribute()), + lambda target: flip.stream_region_target("case", target, [Z, Y, X], Attribute()), lambda window, target, source: flip.inverse("case", window, Attribute()), [Z, Y, X], [Z, Y, X], @@ -116,7 +117,7 @@ def test_stream_orientation_permuted_slab_axis_matches_whole_volume(seed: int) - out_shape = permute.inverse_transform_shape(in_shape, Attribute()) assert out_shape == [Z, Y, X] got = _run_stream( - lambda target: permute.stream_region_target(target, in_shape, Attribute()), + lambda target: permute.stream_region_target("case", target, in_shape, Attribute()), lambda window, target, source: permute.inverse("case", window, Attribute()), in_shape, out_shape, @@ -141,7 +142,7 @@ def test_stream_orientation_canonical_inverse_matches_whole_volume(seed: int) -> assert canonical.inverse_patch_locality(attribute).kind is LocalityKind.ORIENTATION reference = canonical.inverse("case", canonical_volume.clone(), Attribute(attribute)) got = _run_stream( - lambda target: canonical.stream_region_target(target, [Z, Y, X], Attribute(attribute)), + lambda target: canonical.stream_region_target("case", target, [Z, Y, X], Attribute(attribute)), lambda window, target, source: canonical.inverse("case", window, Attribute(attribute)), [Z, Y, X], canonical.inverse_transform_shape([Z, Y, X], attribute), @@ -203,7 +204,7 @@ def test_stream_crop_padding_inverse_matches_whole_volume(seed: int) -> None: reference = padding.inverse("case", padded, Attribute()) assert list(reference.shape[1:]) == out_shape got = _run_stream( - lambda target: padding.stream_region_target(target, in_shape, Attribute()), + lambda target: padding.stream_region_target("case", target, in_shape, Attribute()), lambda window, target, source: window, in_shape, out_shape, @@ -215,8 +216,8 @@ def test_stream_crop_padding_inverse_matches_whole_volume(seed: int) -> None: @pytest.mark.parametrize("seed", range(4)) def test_stream_rescale_nearest_is_byte_identical_to_the_whole_volume_inverse(seed: int) -> None: - # The streamed resample takes its index map from F.interpolate itself, so nearest (uint8) is - # byte-identical to the whole-volume inverse — the exactness the RESCALE gate relies on. + # The streamed inverse runs the code the whole-volume inverse runs, over a region that happens + # to be smaller: byte-identical by construction rather than by agreement. rng = np.random.default_rng(seed) volume = torch.from_numpy(rng.integers(0, 7, size=(C, Z, Y, X)).astype(np.uint8)) resample = ResampleToResolution([1.0, 1.0, 1.0]) @@ -227,12 +228,14 @@ def test_stream_rescale_nearest_is_byte_identical_to_the_whole_volume_inverse(se in_shape = [Z, Y, X] out_shape = resample.inverse_transform_shape(in_shape, attribute) assert out_shape == [12, 9, 7] - scales = [in_shape[k] / out_shape[k] for k in range(3)] reference = resample.inverse("case", volume.clone(), Attribute(attribute)) got = _run_stream( - lambda target: resample.stream_region_target(target, in_shape, Attribute(attribute)), - lambda window, target, source: resample.resample_region( - window, target, [s.start for s in source], scales, in_shape + lambda target: resample.stream_region_target("case", target, in_shape, Attribute(attribute)), + lambda window, target, source: resample.stream_region_inverse( + "case", + window, + RegionContext(tuple(source), tuple(target), tuple(in_shape), tuple(out_shape)), + Attribute(attribute), ), in_shape, out_shape, @@ -244,9 +247,8 @@ def test_stream_rescale_nearest_is_byte_identical_to_the_whole_volume_inverse(se @pytest.mark.parametrize("seed", range(4)) def test_stream_rescale_linear_matches_the_whole_volume_inverse_to_float_rounding(seed: int) -> None: - # A float (linear) rescale is not byte-identical to F.interpolate windowed, but resample_region - # computes the same linear taps, so the streamed inverse matches the whole-volume one to - # ~float-rounding (KONFAI_STREAM_LINEAR_RESAMPLE trades exactly this for a bounded window). + # And a float one too: coordinates are GLOBAL, so a region and the whole volume put the same + # sample in the same place. There is no tolerance to negotiate here any more. rng = np.random.default_rng(seed) volume = torch.from_numpy(rng.standard_normal((C, Z, Y, X)).astype(np.float32)) * 100.0 resample = ResampleToResolution([1.0, 1.0, 1.0]) @@ -256,12 +258,14 @@ def test_stream_rescale_linear_matches_the_whole_volume_inverse_to_float_roundin attribute["Size"] = np.asarray([Z, Y, X]) in_shape = [Z, Y, X] out_shape = resample.inverse_transform_shape(in_shape, attribute) - scales = [in_shape[k] / out_shape[k] for k in range(3)] reference = resample.inverse("case", volume.clone(), Attribute(attribute)) got = _run_stream( - lambda target: resample.stream_region_target(target, in_shape, Attribute(attribute)), - lambda window, target, source: resample.resample_region( - window, target, [s.start for s in source], scales, in_shape + lambda target: resample.stream_region_target("case", target, in_shape, Attribute(attribute)), + lambda window, target, source: resample.stream_region_inverse( + "case", + window, + RegionContext(tuple(source), tuple(target), tuple(in_shape), tuple(out_shape)), + Attribute(attribute), ), in_shape, out_shape, @@ -309,11 +313,13 @@ def test_stream_window_is_bounded_by_the_pull_span() -> None: attribute["Size"] = np.asarray([tall, Y, X]) in_shape = [tall, Y, X] out_shape = [96, Y, X] - scales = [in_shape[k] / out_shape[k] for k in range(3)] stream = SlabRegionStream( - lambda target: resample.stream_region_target(target, in_shape, Attribute(attribute)), - lambda window, target, source: resample.resample_region( - window, target, [s.start for s in source], scales, in_shape + lambda target: resample.stream_region_target("case", target, in_shape, Attribute(attribute)), + lambda window, target, source: resample.stream_region_inverse( + "case", + window, + RegionContext(tuple(source), tuple(target), tuple(in_shape), tuple(out_shape)), + Attribute(attribute), ), in_shape, out_shape, @@ -362,7 +368,7 @@ def test_inverse_locality_defaults_and_overrides() -> None: seeded = Attribute() seeded["Size"] = np.asarray([4, 4, 4]) seeded["Size"] = np.asarray([2, 2, 2]) - assert ResampleToResolution().inverse_patch_locality(seeded).kind is LocalityKind.RESCALE + assert ResampleToResolution().inverse_patch_locality(seeded).kind is LocalityKind.REGRID # Canonical judges the POPPED state: without a stacked direction there is nothing to invert onto. assert Canonical().inverse_patch_locality(empty).kind is LocalityKind.WHOLE_VOLUME @@ -608,7 +614,7 @@ def test_add_layer_streams_a_forward_region_final_transform(tmp_path, monkeypatc def test_add_layer_streams_a_full_geometry_stack_through_the_composed_pipe(tmp_path, monkeypatch) -> None: # The general case the composition exists for: Canonical + ResampleToResolution + Padding forward, - # so the finalize chain carries CROP + RESCALE + ORIENTATION in sequence. With the labelmap cast + # so the finalize chain carries CROP + REGRID + ORIENTATION in sequence. With the labelmap cast # to uint8 before the reduction, the whole stack streams to the sink and must match the # whole-volume path bit for bit. volume = (torch.arange(1 * 6 * 4 * 3).reshape(1, 6, 4, 3) % 5).to(torch.float32) diff --git a/tests/unit/test_transform.py b/tests/unit/test_transform.py index 4384a2b9..468c5689 100644 --- a/tests/unit/test_transform.py +++ b/tests/unit/test_transform.py @@ -292,15 +292,21 @@ def test_padding_after_the_data_keeps_origin(image_attributes): def test_resample_to_resolution_transform_shape_missing_spacing_raises(): - """A tensor without 'Spacing' metadata must surface a TransformError, not fall through.""" + """A density change is meaningless without the density it starts from, so it refuses.""" with pytest.raises(TransformError): ResampleToResolution().transform_shape("group", "case", [10, 10, 10], Attribute()) -def test_resample_to_shape_transform_shape_missing_spacing_raises(): - """ResampleToShape must also raise when 'Spacing' metadata is absent.""" - with pytest.raises(TransformError): - ResampleToShape().transform_shape("group", "case", [10, 10, 10], Attribute()) +def test_resample_to_shape_needs_no_spacing_at_all(): + """A count is a count. Only a DENSITY change needs the density it starts from. + + This used to refuse alongside ``ResampleToResolution``, which cost nothing but told the user to + go and find a geometry for an operation that is a pure resize. With no header the grid is the + identity, and the map degenerates to the size ratio it always was. + """ + assert ResampleToShape(shape=[4, 5, 6]).transform_shape("group", "case", [10, 10, 10], Attribute()) == [4, 5, 6] + resized = ResampleToShape(shape=[4, 5, 6])("case", torch.zeros(1, 10, 10, 10), Attribute()) + assert list(resized.shape[1:]) == [4, 5, 6] def test_resample_to_resolution_transform_shape_dimension_mismatch_message(): @@ -309,7 +315,8 @@ def test_resample_to_resolution_transform_shape_dimension_mismatch_message(): attributes["Spacing"] = np.asarray([1.0, 1.0], dtype=np.float64) with pytest.raises(TransformError) as excinfo: ResampleToResolution(spacing=[1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], attributes) - assert "shape=[10, 10, 10]" in str(excinfo.value) + message = str(excinfo.value) + assert "case 'case'" in message and "3-dimensional grid" in message def test_resample_to_shape_transform_shape_dimension_mismatch_message(): @@ -319,19 +326,21 @@ def test_resample_to_shape_transform_shape_dimension_mismatch_message(): with pytest.raises(TransformError) as excinfo: ResampleToShape(shape=[4, 4]).transform_shape("group", "case", [10, 10, 10], attributes) message = str(excinfo.value) - assert "shape=[10, 10, 10]" in message - assert "target_shape" in message + assert "shape of 2 value(s)" in message + assert "3 spatial axis/axes" in message def test_resample_to_shape_does_not_mutate_config(): """#9 transform_shape must not write resolved dims back into the shared instance config.""" resampler = ResampleToShape(shape=[0, 16, 16]) - before = resampler.shape.clone() attributes = Attribute() attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) - out = resampler.transform_shape("CT", "case", [8, 16, 16], attributes) - assert out[0] == 8 # sentinel 0 resolved to the input dim for this call - assert torch.equal(resampler.shape, before), "self.shape must stay [0, 16, 16] for the next case" + + first = resampler.transform_shape("CT", "case_a", [8, 16, 16], Attribute(attributes)) + assert first[0] == 8 # sentinel 0 resolved to the input dim for this call + # The sentinel has to survive it: the next case is a different volume, and a stage is shared. + second = resampler.transform_shape("CT", "case_b", [11, 16, 16], Attribute(attributes)) + assert second[0] == 11 def test_resample_to_shape_inverse_without_spacing_metadata(): diff --git a/tests/unit/test_transform_bound.py b/tests/unit/test_transform_bound.py new file mode 100644 index 00000000..42da3d3f --- /dev/null +++ b/tests/unit/test_transform_bound.py @@ -0,0 +1,255 @@ +# 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 + +"""What a decoded transform's bound must satisfy, against SimpleITK. + +Two properties, and both matter. CONTAINMENT: the true map never leaves the bound — a bound that +is short reads a source window the resample then samples outside of, which returns background +rather than failing. NON-VACUITY: the bound is not the whole world — a bound that contains +everything contains the truth and buys nothing. +""" + +import numpy as np +import pytest +from konfai.data.geometry import AffineStage, DisplacementStage, bound_of +from konfai.utils.errors import TransformError + +sitk = pytest.importorskip("SimpleITK") + +from konfai.utils.ITK import decode_transform_stages, invert_stages # noqa: E402 + +SIZE = (32, 40, 48) + + +def _image(oblique: bool = True) -> "sitk.Image": + image = sitk.GetImageFromArray(np.zeros(SIZE, np.float32)) + image.SetSpacing((0.8, 1.2, 1.5)) + image.SetOrigin((10.0, -5.0, 2.0)) + if oblique: + a, b = np.deg2rad(20.0), np.deg2rad(15.0) + rz = np.array([[np.cos(a), -np.sin(a), 0.0], [np.sin(a), np.cos(a), 0.0], [0.0, 0.0, 1.0]]) + ry = np.array([[np.cos(b), 0.0, np.sin(b)], [0.0, 1.0, 0.0], [-np.sin(b), 0.0, np.cos(b)]]) + image.SetDirection(tuple((rz @ ry).ravel())) + return image + + +def _euler(image: "sitk.Image") -> "sitk.Transform": + transform = sitk.Euler3DTransform() + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetRotation(0.11, -0.2, 0.31) + transform.SetTranslation((3.0, -2.0, 1.0)) + return transform + + +def _affine(image: "sitk.Image") -> "sitk.Transform": + transform = sitk.AffineTransform(3) + transform.SetCenter(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + transform.SetMatrix(np.array([[1.1, 0.05, 0.0], [0.0, 0.9, 0.07], [0.02, 0.0, 1.2]]).ravel()) + transform.SetTranslation((4.0, -3.0, 2.0)) + return transform + + +def _bspline(image: "sitk.Image", mesh: int = 6, amplitude: float = 9.0) -> "sitk.Transform": + transform = sitk.BSplineTransformInitializer(image, [mesh] * 3, 3) + size = np.asarray(transform.GetParameters()).size + transform.SetParameters(list(np.random.RandomState(1).uniform(-amplitude, amplitude, size))) + return transform + + +def _field(image: "sitk.Image") -> "sitk.Transform": + filt = sitk.TransformToDisplacementFieldFilter() + filt.SetReferenceImage(image) + return sitk.DisplacementFieldTransform(sitk.Cast(filt.Execute(_bspline(image)), sitk.sitkVectorFloat64)) + + +def _world_points(image: "sitk.Image", count: int = 4000, overshoot: float = 6.0) -> np.ndarray: + """Points across the grid AND past its edges — the bound must hold everywhere, not just inside.""" + rng = np.random.RandomState(0) + index = np.stack([rng.uniform(-overshoot, extent + overshoot, count) for extent in image.GetSize()], axis=1) + return np.array([image.TransformContinuousIndexToPhysicalPoint(list(point)) for point in index]) + + +def _cases(image: "sitk.Image") -> list[tuple[str, "sitk.Transform"]]: + return [ + ("euler", _euler(image)), + ("affine", _affine(image)), + ("bspline coarse", _bspline(image, mesh=4)), + ("bspline fine", _bspline(image, mesh=12)), + ("field", _field(image)), + ("composite affine+bspline", sitk.CompositeTransform([_affine(image), _bspline(image)])), + ("composite affine+euler", sitk.CompositeTransform([_affine(image), _euler(image)])), + ] + + +@pytest.mark.parametrize("oblique", [False, True], ids=["axis-aligned", "oblique"]) +def test_the_bound_contains_the_map_it_bounds(oblique: bool): + image = _image(oblique) + world = _world_points(image) + for label, transform in _cases(image): + bound = bound_of(decode_transform_stages(transform), 3) + truth = np.array([transform.TransformPoint(point) for point in world]) + excess = np.abs(truth - bound.affine.apply(world)) - bound.residual_xyz + assert excess.max() <= 1e-9, f"{label}: the bound is short by {excess.max():.4g} mm" + + +@pytest.mark.parametrize("oblique", [False, True], ids=["axis-aligned", "oblique"]) +def test_the_bound_is_not_vacuous(oblique: bool): + # Containment alone is satisfied by an infinite box. The residual must stay within a small + # multiple of the displacement actually reached, or the halo it sizes is the whole volume. + image = _image(oblique) + world = _world_points(image) + for label, transform in _cases(image): + bound = bound_of(decode_transform_stages(transform), 3) + if not bound.residual_xyz.any(): + continue + truth = np.array([transform.TransformPoint(point) for point in world]) + reached = np.abs(truth - bound.affine.apply(world)).max(axis=0) + assert (bound.residual_xyz <= 4.0 * np.maximum(reached, 1e-6)).all(), f"{label}: bound far too loose" + + +def test_a_linear_transform_is_bounded_exactly(): + image = _image() + world = _world_points(image) + for transform in (_euler(image), _affine(image), sitk.TranslationTransform(3, (5.0, -1.0, 2.0))): + bound = bound_of(decode_transform_stages(transform), 3) + np.testing.assert_array_equal(bound.residual_xyz, np.zeros(3)) + truth = np.array([transform.TransformPoint(point) for point in world]) + np.testing.assert_allclose(bound.affine.apply(world), truth, rtol=0.0, atol=1e-9) + + +def test_probing_the_affine_part_of_a_bspline_is_worse_than_the_structural_one(): + # Why the affine part is read structurally and never by finite differences. Probing a NON-linear + # map measures a local gradient and extrapolates it across the whole grid: around an interior + # point of this spline that implies a residual of ~19 mm where the sup-norm of the coefficients + # says 9. A design that probes pays for a halo twice as wide and still has no theorem behind it. + image = _image() + transform = _bspline(image) + world = _world_points(image) + structural = bound_of(decode_transform_stages(transform), 3) + + centre = np.array(image.TransformContinuousIndexToPhysicalPoint([(s - 1) / 2 for s in image.GetSize()])) + base = np.array(transform.TransformPoint(tuple(centre))) + columns = [np.array(transform.TransformPoint(tuple(centre + np.eye(3)[k]))) - base for k in range(3)] + probed = np.stack(columns, axis=1) + truth = np.array([transform.TransformPoint(point) for point in world]) + probed_residual = np.abs(truth - (world @ probed.T + (base - probed @ centre))).max(axis=0) + assert (probed_residual > structural.residual_xyz).any() + + +def test_the_bound_is_a_theorem_and_the_sampled_maximum_is_not(): + # The sup-norm bound must dominate what any sample can reach — that is the whole claim. A design + # that sized its halo from a dense sample would be under a bound it never proved. + image = _image() + transform = _bspline(image) + world = _world_points(image, count=20000) + bound = bound_of(decode_transform_stages(transform), 3) + truth = np.array([transform.TransformPoint(point) for point in world]) + assert np.all(np.abs(truth - world).max(axis=0) <= bound.residual_xyz + 1e-9) + + +class TestRefusals: + def test_a_transform_that_decomposes_into_nothing_is_refused_by_name(self): + class Opaque: + def GetDimension(self): + return 3 + + def GetName(self): + return "MadeUpTransform" + + def IsLinear(self): + return False + + with pytest.raises(TransformError, match="MadeUpTransform"): + decode_transform_stages(Opaque()) + + def test_a_non_finite_field_is_refused_rather_than_bounded(self): + image = _image() + filt = sitk.TransformToDisplacementFieldFilter() + filt.SetReferenceImage(image) + array = sitk.GetArrayFromImage(filt.Execute(_bspline(image))) + array[0, 0, 0, 0] = np.nan + field = sitk.GetImageFromArray(array, isVector=True) + field.CopyInformation(image) + transform = sitk.DisplacementFieldTransform(sitk.Cast(field, sitk.sitkVectorFloat64)) + with pytest.raises(TransformError, match="non-finite"): + decode_transform_stages(transform) + + +class TestCompositeOrder: + def test_the_stages_are_decoded_in_application_order(self): + # SimpleITK applies a CompositeTransform's list in REVERSE (last added runs first), while + # GetNthTransform(0) is the first added. Folding a bound in list order composes the wrong + # way round; the decoder normalizes it, and this is what pins that. + translate = sitk.TranslationTransform(3, (100.0, 0.0, 0.0)) + scale = sitk.ScaleTransform(3, (2.0, 2.0, 2.0)) + composite = sitk.CompositeTransform([translate, scale]) + point = np.array([1.0, 1.0, 1.0]) + np.testing.assert_allclose(composite.TransformPoint(tuple(point)), [102.0, 2.0, 2.0]) + + stages = decode_transform_stages(composite) + folded = bound_of(stages, 3) + np.testing.assert_allclose(folded.affine.apply(point), [102.0, 2.0, 2.0], rtol=0.0, atol=1e-9) + + def test_the_residual_is_transported_through_an_outer_scale(self): + image = _image(oblique=False) + spline = _bspline(image, mesh=5, amplitude=5.0) + alone = bound_of(decode_transform_stages(spline), 3) + scaled = bound_of( + decode_transform_stages(sitk.CompositeTransform([sitk.ScaleTransform(3, (3.0,) * 3), spline])), 3 + ) + # The spline runs first, then the scale: its reach is multiplied by 3, not left alone. + np.testing.assert_allclose(scaled.residual_xyz, 3.0 * alone.residual_xyz, rtol=1e-9) + + +class TestInverse: + def test_an_all_affine_map_inverts_algebraically(self): + image = _image() + transform = sitk.CompositeTransform([_affine(image), _euler(image)]) + inverse = invert_stages(decode_transform_stages(transform), 3) + assert inverse is not None + forward = bound_of(decode_transform_stages(transform), 3).affine + back = bound_of(inverse, 3).affine + points = _world_points(image, count=200) + np.testing.assert_allclose(back.apply(forward.apply(points)), points, rtol=0.0, atol=1e-6) + + @pytest.mark.parametrize("factory", [_bspline, _field], ids=["bspline", "field"]) + def test_a_non_affine_map_has_no_algebraic_inverse(self, factory): + assert invert_stages(decode_transform_stages(factory(_image())), 3) is None + + +class TestDisplacementStage: + def test_a_bspline_decodes_to_its_coefficient_grid(self): + image = _image() + stages = decode_transform_stages(_bspline(image, mesh=7)) + assert len(stages) == 1 + stage = stages[0] + assert isinstance(stage, DisplacementStage) + assert stage.order == 3 + # The bound IS the sup-norm of the coefficients, per component. + np.testing.assert_allclose(stage.bound_xyz, np.abs(stage.values.reshape(3, -1)).max(axis=1)) + + def test_a_field_decodes_to_order_one_on_its_own_grid(self): + image = _image() + stages = decode_transform_stages(_field(image)) + stage = stages[0] + assert isinstance(stage, DisplacementStage) + assert stage.order == 1 + assert stage.grid.size_zyx == SIZE + + def test_an_affine_stage_carries_no_residual(self): + stages = decode_transform_stages(_euler(_image())) + assert isinstance(stages[0], AffineStage) + np.testing.assert_array_equal(stages[0].bound().residual_xyz, np.zeros(3)) diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index 5a7333bf..252bb59b 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -38,6 +38,7 @@ import numpy as np import pytest +import SimpleITK as sitk import torch from konfai.data import augmentation as augmentation_module from konfai.data import transform as transform_module @@ -116,14 +117,16 @@ # float32 ulps away. Data-dependent: this fixture happens to agree exactly, a smooth field showed # 1.5e-8 (0.13 ulp), so the bound is stated rather than observed. _STAT_ATOL = 8 * float(np.finfo(np.float32).eps) -# - the streamed resample (trilinear only): it gathers the same source samples, but computes the -# interpolation weights from coordinates expressed in the read sub-region's frame rather than the -# whole volume's. Both round to float32, so a weight lands ~ulp(coordinate) off and the interpolated -# voxel lands `neighbour gap * ulp(coordinate)` off -- the deviation scales with the local GRADIENT, -# not with the voxel's own magnitude. The fixture's gap is its 2*_PEAK bone/air step, which puts the -# bound at ulps of _PEAK; 64 of them is ~8x the measured max (2.3e-4, i.e. 7.5 ulp) and stays far -# below one part per million of the range. Nearest (uint8) uses no weights and stays exact. -_RESCALE_ATOL = 64 * float(np.spacing(np.float32(_PEAK))) +# - a streamed regrid through a map that does NOT factorise (a field, a rotation): its blend goes to +# grid_sample, which takes normalised coordinates and so expresses them in the read sub-region's +# frame rather than the whole volume's. Both round to float32, so a weight lands ~ulp(coordinate) +# off and the interpolated voxel lands `neighbour gap * ulp(coordinate)` off -- the deviation +# scales with the local GRADIENT, not with the voxel's own magnitude. The fixture's gap is its +# 2*_PEAK bone/air step, which puts the bound at ulps of _PEAK; 64 of them is ~8x the measured max +# and stays far below one part per million of the range. A map that DOES factorise is read one axis +# at a time on global coordinates and stays bit-identical, which is why those cases carry no atol +# at all; nearest uses no weights and is exact either way. +_REGRID_ATOL = 64 * float(np.spacing(np.float32(_PEAK))) # An integer volume truncates the interpolation, so a sub-ulp disagreement that straddles an integer # boundary becomes a whole least-significant bit. 1 LSB is the tightest bound that can hold: the # alternative would be bit-exact agreement between two different float coordinate frames. @@ -168,12 +171,12 @@ class _Case: "Percentage": [_Case(Percentage(100.0))], # The defaults ([1, 1, 1] mm / [100, 256, 256]) would be a no-op resample and a 6.5M-voxel upsample. "ResampleToResolution": [ - _Case(ResampleToResolution([2.0, 1.0, 3.0]), atol=_RESCALE_ATOL), + _Case(ResampleToResolution([2.0, 1.0, 3.0])), # factorises: bit-identical _Case(ResampleToResolution([2.0, 1.0, 3.0]), group="Int16", atol=_LSB_ATOL), # uint8 resamples by nearest neighbour: no interpolation weights, so no rounding to disagree on. _Case(ResampleToResolution([2.0, 1.0, 3.0]), group="Labels"), ], - "ResampleToShape": [_Case(ResampleToShape([12, 8, 14]), atol=_RESCALE_ATOL)], + "ResampleToShape": [_Case(ResampleToShape([12, 8, 14]))], # factorises: bit-identical # 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 @@ -185,11 +188,18 @@ class _Case: # 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. + # Through a field the map does not factorise, so the blend goes to grid_sample -- which + # normalises by the extent it is handed, and a patch is handed a window. That is the one + # place a streamed answer is not bit-identical to the whole-volume one, and the atol says so. _Case( - ResampleToReference(entry=_CASE_NAME, group="Reference", field_group="Field", max_displacement=_FIELD_BOUND) + ResampleToReference( + entry=_CASE_NAME, group="Reference", field_group="Field", max_displacement=_FIELD_BOUND + ), + atol=_REGRID_ATOL, ), ], - "ResampleTransform": [_Case(ResampleTransform({"transform": True}))], + # A stored map never factorises, so this is the grid_sample path (see _REGRID_ATOL). + "ResampleTransform": [_Case(ResampleTransform({"transform": True}), atol=_REGRID_ATOL)], "Save": [_Case(Save("Dataset"))], # Warp needs a field on disk to run, which this registry cannot build: with no declared # displacement it declares WHOLE_VOLUME, so it stays out of the equivalence sweep below. Its @@ -295,9 +305,19 @@ def _attributes(group: str) -> Attribute: @pytest.fixture(scope="session") def dataset(tmp_path_factory: pytest.TempPathFactory) -> Dataset: """A real on-disk dataset, in the same format (mha) and channel-first layout a run reads.""" - dataset = Dataset(tmp_path_factory.mktemp("workspace") / "Dataset", "mha") + root = tmp_path_factory.mktemp("workspace") / "Dataset" + dataset = Dataset(root, "mha") for group, volume in _volumes().items(): dataset.write(group, _CASE_NAME, volume, _attributes(group)) + # A stored transform, as a group of its own: KonfAI reads one from `/.itk.txt`, + # which is what SimpleITK writes. Without it the resample-through-a-stored-map case has nothing + # to apply and drops out of the sweep, which is the one kind of hole a registry cannot show. + stored = sitk.Euler3DTransform() + stored.SetCenter((3.0, 5.0, 11.0)) + stored.SetRotation(0.05, -0.03, 0.08) + stored.SetTranslation((0.4, -0.6, 1.1)) + (root / _CASE_NAME).mkdir(parents=True, exist_ok=True) + sitk.WriteTransform(stored, str(root / _CASE_NAME / "transform.itk.txt")) return dataset @@ -501,9 +521,9 @@ def test_a_streamed_region_records_the_whole_volume_geometry(dataset: Dataset, g # A HALO draw is sampled by grid_sample from coordinates expressed in the halo'd read extent's frame # rather than the whole volume's: the same disagreement, for the same reason and with the same -# gradient- and coordinate-scaling, that _RESCALE_ATOL bounds for the streamed resample. It is bitwise +# gradient- and coordinate-scaling, that _REGRID_ATOL bounds for the streamed resample. It is bitwise # on neither, and grows with the extent -- a 160^3 case at patch 64 lands at 2e-5 of its range. -_AUGMENTATION_ATOL = _RESCALE_ATOL +_AUGMENTATION_ATOL = _REGRID_ATOL @dataclass(frozen=True) diff --git a/tests/unit/test_transformer_workflow.py b/tests/unit/test_transformer_workflow.py index 3d47ec65..b0411ec5 100644 --- a/tests/unit/test_transformer_workflow.py +++ b/tests/unit/test_transformer_workflow.py @@ -396,8 +396,9 @@ def spy(destination, group, shape, dtype, attributes): _RESAMPLED_THEN_REFERENCED = """\ - ResampleToResolution: + Resample: spacing: [2.0, 2.0, 2.0] + align: origin ResampleToReference: entry: CASE_000 group: CT @@ -409,9 +410,14 @@ def spy(destination, group, shape, dtype, attributes): 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. + The resample takes 12x10x8 down to 12x7x2 and keeps voxel zero where it is, so the case's far + edge falls short of the reference's and part of the output will be fill. Asked about the case as + STORED it covers all of it and says nothing -- and the plan would stay silent about that fill. + + (With ``align: extent``, the default, the box is preserved and the answer is honestly 100%. + 67.5% is also exactly what the old ``ResampleToResolution`` recorded here -- because its header + said origin-aligned while its data was extent-aligned, which is the mismatch this stage no + longer has.) """ _write_source(tmp_path) _write_config(tmp_path, _RESAMPLED_THEN_REFERENCED.format(out=tmp_path / "out")) diff --git a/tests/unit/test_warp.py b/tests/unit/test_warp.py index 5d259e67..b047c396 100644 --- a/tests/unit/test_warp.py +++ b/tests/unit/test_warp.py @@ -81,25 +81,55 @@ def _manager(source: Dataset, transforms: list) -> DatasetManager: ) -def test_declares_a_halo_sized_by_the_declared_displacement() -> None: - warp = Warp(field="./x:h5", group="DVF", max_displacement=4.0) - locality = warp.patch_locality(_attributes()) - # Spacing is (x=2, y=1, z=1) so in array order (z, y, x) it is (1, 1, 2): 4 um is 4, 4 and 2 voxels. - assert locality.kind is LocalityKind.HALO - assert locality.halo == (4, 4, 2) +def _recorded(warp: Warp, attribute: Attribute | None = None, shape: tuple[int, ...] = (10, 12, 14)) -> Warp: + """A stage that has met its case — which is when a region can be asked about at all.""" + warp.transform_shape("CT", "CASE_000", list(shape), attribute if attribute is not None else _attributes()) + return warp + + +def test_the_source_region_is_the_target_grown_by_the_declared_displacement() -> None: + """A warp is a regrid onto the case's own grid, and its window is the bound in voxels. + + Spacing is (x=2, y=1, z=1), so in array order (z, y, x) 4 um of displacement is 4, 4 and 2 + voxels -- plus the one voxel the linear taps reach. Declared as REGRID and not HALO because the + window is derived from the case's GEOMETRY: see the oblique case below, which a per-axis halo + cannot express at all. + """ + warp = _recorded(Warp(field="./x:h5", group="DVF", max_displacement=4.0)) + assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID + target = (slice(4, 6), slice(4, 6), slice(4, 6)) + window = warp.stream_region_source("CASE_000", target, [10, 12, 14], _attributes()) + + # The rule, written out: the region's OUTER faces (start - 0.5 .. stop - 0.5) in world units, + # grown by the declared 4 um, back to indices, floor/ceil, one voxel of margin for the taps. + extents, per_voxel = (10, 12, 14), (1.0, 1.0, 2.0) # array order (z, y, x) + expected = [] + for axis, extent in enumerate(extents): + reach = 4.0 / per_voxel[axis] + low, high = 4 - 0.5 - reach, 6 - 0.5 + reach + expected.append((max(0, int(np.floor(low)) - 1), min(extent, int(np.ceil(high)) + 2))) + assert [(part.start, part.stop) for part in window] == expected + + +def test_an_oblique_case_grows_its_window_on_every_axis() -> None: + """The bug a per-axis halo hid: a displacement along x reaches into y and z when the axes turn. + + ``Warp`` used to convert a world bound to a halo per ARRAY axis, which silently assumed the + direction cosines were the identity -- on a turned case the window was short on the axes the + displacement actually reached, and a short window returns the border value rather than raising. + """ + turned = _attributes() + angle = np.deg2rad(35.0) + cos, sin = float(np.cos(angle)), float(np.sin(angle)) + turned["Direction"] = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]).reshape(-1) -def test_without_a_declared_bound_it_refuses_to_stream() -> None: - """The safety net: an undeclared reach is an unbounded read, so the whole volume it is.""" - undeclared = Warp(field="./x:h5", group="DVF").patch_locality(_attributes()) - assert undeclared.kind is LocalityKind.WHOLE_VOLUME - # And it says which of the two is missing: a Warp that silently costs the whole volume still - # produces the right result, so the reason is the only thing that shows. - assert undeclared.reason is not None and "max_displacement" in undeclared.reason + warp = _recorded(Warp(field="./x:h5", group="DVF", max_displacement=4.0), turned) + target = (slice(5, 6), slice(5, 6), slice(5, 6)) + window = warp.stream_region_source("CASE_000", target, [10, 12, 14], turned) - no_geometry = Warp(field="./x:h5", group="DVF", max_displacement=4.0).patch_locality(Attribute()) - assert no_geometry.kind is LocalityKind.WHOLE_VOLUME - assert no_geometry.reason is not None and "Spacing" in no_geometry.reason + widths = [part.stop - part.start for part in window] + assert all(width > 1 for width in widths), f"a turned case reaches on every axis, got {widths}" @_needs_rfc5 @@ -122,9 +152,13 @@ def test_auto_reads_the_bound_the_fields_recorded_when_they_were_written(tmp_pat warp = Warp(field=f"{tmp_path / 'dvf'}:omezarr", group="DVF", max_displacement="auto") locality = warp.patch_locality(_attributes()) - # The cohort's bound is (x=1.0, y=6.0, z=3.0); spacing in array order (z, y, x) is (1, 1, 2). - assert locality.kind is LocalityKind.HALO - assert locality.halo == (3, 6, 1) + # The cohort's bound is (x=1.0, y=6.0, z=3.0); spacing in array order (z, y, x) is (1, 1, 2), so + # the window grows by 3, 6 and 1 voxels (plus the linear taps' one) around its target. + assert locality.kind is LocalityKind.REGRID + window = _recorded(warp).stream_region_source("CASE_000", (slice(4, 6),) * 3, [10, 12, 14], _attributes()) + reaches = (3.0, 6.0, 0.5) # array order (z, y, x): the bound divided by that axis's spacing + starts = [max(0, int(np.floor(4 - 0.5 - reach)) - 1) for reach in reaches] + assert [part.start for part in window] == starts def test_auto_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> None: @@ -200,6 +234,7 @@ def test_a_field_beyond_the_declared_bound_raises(tmp_path: Path) -> None: _source, _fields, volume = _fixture(tmp_path, shift_um=(0.0, 0.0, 9.0)) warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement=1.0) + _recorded(warp) with pytest.raises(TransformError, match=r"on component 2, beyond the 1\.000"): whole = (slice(0, 10), slice(0, 12), slice(0, 14)) warp.stream_region( diff --git a/tests/unit/test_write_pyramid_and_field_bound.py b/tests/unit/test_write_pyramid_and_field_bound.py index 34a22473..7dc91cab 100644 --- a/tests/unit/test_write_pyramid_and_field_bound.py +++ b/tests/unit/test_write_pyramid_and_field_bound.py @@ -191,27 +191,51 @@ def test_a_field_records_its_own_bound_on_both_write_paths(tmp_path): @_needs_rfc5 -def test_the_recorded_bound_turns_into_a_per_axis_halo_under_anisotropy(tmp_path) -> None: +def tmp_field_store(field: np.ndarray) -> Path: + import tempfile + + root = Path(tempfile.mkdtemp()) + store = root / "fields" / "case" / "DVF.ome.zarr" + store.parent.mkdir(parents=True) + write_ome_zarr(store, field, spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0), displacement_field=True) + return root / "fields" + + +def test_the_recorded_bound_reaches_each_axis_by_its_own_spacing_under_anisotropy() -> None: """What the bound is FOR, read by the stage that consumes it. Component ``i`` of a displacement field is world axis (x, y, z)[i], while array axes are - (z, y, x). So a halo in array order reads the components reversed, each against its own spacing. - Getting the pairing wrong is a warp that raises nothing and reads the wrong neighbourhood. + (z, y, x). So the reach of a region on each array axis reads the components reversed, each + against its own spacing. Getting the pairing wrong is a warp that raises nothing and reads the + wrong neighbourhood -- and under this anisotropy the three numbers are far enough apart (3, 22 + and 31 voxels) that any permutation of them is visible. """ from konfai.data.transform import LocalityKind, Warp field = np.zeros((3, 8, 8, 8), dtype=np.float32) field[0, 4, 4, 4], field[1, 2, 2, 2], field[2, 1, 1, 1] = 917.5, -640.25, 96.0 - store = tmp_path / "fields" / "case" / "DVF.ome.zarr" - store.parent.mkdir(parents=True) - write_ome_zarr(store, field, spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0), displacement_field=True) + store = tmp_field_store(field) - warp = Warp(field=f"{tmp_path / 'fields'}:omezarr", group="DVF", max_displacement="auto") + warp = Warp(field=f"{store}:omezarr", group="DVF", max_displacement="auto") attribute = Attribute() attribute["Spacing"] = np.array([30.08, 30.08, 40.0]) # stored (x, y, z) + attribute["Origin"] = np.zeros(3) + attribute["Direction"] = np.eye(3).reshape(-1) + + assert warp.patch_locality(attribute).kind is LocalityKind.REGRID - locality = warp.patch_locality(attribute) + shape = [64, 128, 128] + warp.transform_shape("CT", "CASE_000", shape, attribute) + target = tuple(slice(30, 32) for _ in shape) + window = warp.stream_region_source("CASE_000", target, shape, attribute) # z takes the z component (96 um over a 40 um voxel), x the x component (917.5 over 30.08). - assert locality.kind is LocalityKind.HALO - assert locality.halo == (3, 22, 31) + per_axis = [(96.0, 40.0), (640.25, 30.08), (917.5, 30.08)] # array order (z, y, x) + expected = [ + ( + max(0, int(np.floor(30 - 0.5 - reach / spacing)) - 1), + min(extent, int(np.ceil(32 - 0.5 + reach / spacing)) + 2), + ) + for (reach, spacing), extent in zip(per_axis, shape, strict=True) + ] + assert [(part.start, part.stop) for part in window] == expected From f98bd4f20da6d70c1949ea6ebac2a29662a79290 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 23:18:04 +0200 Subject: [PATCH 02/39] fix(data): the correctness pass over the streamed resample path - judge every landing fold on the state the stages before it left - warp the end plane of a B-spline's valid region as ITK warps it - judge coverage through the declared map before refusing a case as disjoint - blend through float32 coordinates whatever the payload dtype - refuse at plan time the map neither route can apply - replay a copy's regions on that copy's own recorded grids - clip to the case's seeded statistic, not the region's own --- konfai/data/patching.py | 70 ++++++++--- konfai/data/sampling.py | 26 +++- konfai/data/transform.py | 108 +++++++++++++--- tests/unit/test_case_expansion.py | 35 ++++++ tests/unit/test_resample_sampler_rules.py | 20 +++ tests/unit/test_resample_transform.py | 86 ++++++++++--- tests/unit/test_streamed_read_dispatcher.py | 133 ++++++++++++++++++++ tests/unit/test_transform_bound.py | 24 ++++ 8 files changed, 444 insertions(+), 58 deletions(-) diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 1808db96..4fcbd963 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -1356,10 +1356,13 @@ def __init__( # The chain around its Expand: pre runs once per case, post once per copy. Without a marker # the split is (everything, None, []) and every fold below reduces to the plain per-case chain. self._expand_pre, self._expand, self._expand_post = split_expand(transforms) + # The landing fold, on a working state: the run re-makes every transition itself, so the + # case baseline the walks and replays start from must not carry them twice. + folding = Attribute(cache_attribute) for transform_function in self._expand_pre: - _shape = _spatial(transform_function.transform_shape(self.group_src, self.name, _shape, cache_attribute)) + _shape = self._fold_case_state(transform_function, _shape, folding) + self._adopt_case_facts(folding, cache_attribute) # The grid and case state at the Expand point: what the first per-copy stage is handed. - # Snapshots (Attribute copies deeply) -- the live attribute keeps evolving through post. self._shape_at_expand = list(_shape) self._attributes_at_expand = Attribute(cache_attribute) # The un-augmented landing of the per-copy tail. A draw is the identity here, because copy 0 @@ -1367,7 +1370,8 @@ def __init__( for transform_function in self._expand_post: if _is_draw(transform_function): continue - _shape = _spatial(transform_function.transform_shape(self.group_src, self.name, _shape, cache_attribute)) + _shape = self._fold_case_state(transform_function, _shape, folding) + self._adopt_case_facts(folding, cache_attribute) self.patch = ( DatasetPatch( @@ -1461,6 +1465,9 @@ def _draw_expand_copies(self, reset_state: bool) -> None: assert expand is not None # nosec B101 - the caller checked shapes = [list(self._shape_at_expand) for _ in range(expand.nb)] attributes = [copy.deepcopy(self._attributes_at_expand) for _ in range(expand.nb)] + # The copies' walk states, apart from the baselines above: the landing fold evolves the + # geometry, and a streamed replay must start from the case as stored. + foldings = [Attribute(attribute) for attribute in attributes] drawn: dict[str, int] = {} for stage in self._expand_post: if _is_draw(stage): @@ -1475,10 +1482,7 @@ def _draw_expand_copies(self, reset_state: bool) -> None: shapes = stage.state_init(self.index, shapes, attributes) continue for index in range(expand.nb): - shapes[index] = [ - int(extent) - for extent in stage.transform_shape(self.group_src, self.name, shapes[index], attributes[index]) - ] + shapes[index] = self._fold_case_state(stage, shapes[index], foldings[index]) for index in range(expand.nb): self.cache_attributes.append(attributes[index]) self.shapes.append(list(shapes[index])) @@ -1826,14 +1830,34 @@ def _stage_out_shape(self, stage: Stage, shape: list[int], attribute: Attribute) The one dispatch between the two Stage species' shape vocabularies: a ``Transform`` restates its fold as ``transform_shape``, an :class:`AugmentedStage` as its draw's ``stream_shape``. - ``attribute`` is used as handed over: a caller that wants the geometry rewrites a transform - makes along the way (the write probe) passes the live state, one that must not be touched - (the planner) passes a copy. + Shape only — the geometry transition is :meth:`_fold_case_state`'s half. """ if isinstance(stage, Transform): return [int(e) for e in stage.transform_shape(self.group_src, self.name, list(shape), attribute)] return [int(e) for e in cast(AugmentedStage, stage).stream_shape(list(shape))] + def _fold_case_state(self, stage: Stage, shape: list[int], attribute: Attribute) -> list[int]: + """Fold one stage over the evolving case state: the shape through its map, the geometry + through its stated transition — the idiom :meth:`_plan_read_stage` runs per region stage. + + Every landing fold goes through here, so a stage is judged on the state the stages before it + left rather than on the stored header — a ``Resample`` behind a ``Canonical`` records the + reoriented grid, and a second ``Resample`` sees the first one's spacing. + """ + out = self._stage_out_shape(stage, shape, attribute) + if isinstance(stage, Transform): + stage.write_stream_cache_attribute(attribute, list(shape)) + return out + + @staticmethod + def _adopt_case_facts(folding: Attribute, case: Attribute) -> None: + """Keep what a landing fold computed about the CASE — Crop's content-derived box — off its + walk state. The geometry the fold evolved is the walk's own (the run re-makes those + transitions); the box is expensive, immutable per case, and read by every later fold, the + streamed replays and the run itself.""" + if "box" in folding and "box" not in case: + case["box"] = folding["box"] + def chain_stages(self, a: int = 0) -> list[Stage]: """The ordered stages copy ``a`` is made of — the one definition of what a copy IS. @@ -1856,7 +1880,7 @@ def write_targets(self, a: int = 0) -> list[tuple[Save, list[int], Attribute]]: attributes = Attribute(self.cache_attributes_bak[0]) targets: list[tuple[Save, list[int], Attribute]] = [] for stage in self.chain_stages(a): - spatial = self._stage_out_shape(stage, spatial, attributes) + spatial = self._fold_case_state(stage, spatial, attributes) if isinstance(stage, Save): targets.append((stage, list(spatial), Attribute(attributes))) return targets @@ -2019,7 +2043,7 @@ def _plan_save_sweep( landing = [int(extent) for extent in source_shape[1:]] probe = Attribute(base_attributes) for stage in segment: - landing = self._stage_out_shape(stage, landing, probe) + landing = self._fold_case_state(stage, landing, probe) planning = Attribute(base_attributes) streamable, stage_plans, evolved, refusal = self._plan_stream_region( 0, @@ -2215,9 +2239,7 @@ def peak_case_bytes(self) -> int: for stage in self.transforms: if not isinstance(stage, Transform): continue - source = list(spatial) - spatial = [int(extent) for extent in stage.transform_shape(self.group_src, self.name, source, attributes)] - stage.write_stream_cache_attribute(attributes, source) + spatial = self._fold_case_state(stage, list(spatial), attributes) peak = max(peak, channels * int(np.prod(spatial, dtype=np.int64))) return peak * CASE_ELEMENT_BYTES @@ -2759,6 +2781,8 @@ def _get_streamed_region_data( """Patch-native region chain: one target patch replayed through the composed region plans, padded back to ``patch_size`` like the whole-volume path (see ``_replay_streamed_region``, which a Save sweep drives with slab targets instead of patch targets).""" + if self._expand is not None: + self._refold_copy_records(a, stream_source) target_slices = tuple(self.patch.get_patch_slices(a)[index]) # Each patch re-runs the chain from the state the whole-volume pass started from: the case as # stored (plus planned stats), never the live attribute -- that one carries the chain's own @@ -2775,6 +2799,22 @@ def _get_streamed_region_data( self._persist_stream_attributes(a, cache_attribute, keys_before) return tensor, cache_attribute + def _refold_copy_records(self, a: int, stream_source: _PatchStreamSource) -> None: + """Re-fold copy ``a``'s chain state before replaying a region of it. + + A stage keys its per-case records by the CASE name — a stored transform is looked up by + it — so the copies of an Expand share one key and the last WALK's records win. The write + sweeps re-plan before sweeping and the whole-volume path re-records at call time; the + patch replay is the consumer left over, and reading two copies interleaved would otherwise + hand one copy the other's grids. Headers only — no voxel is read. + """ + if not stream_source.stages: + return + shape = list(stream_source.stage_plans[0].in_shape) + state = Attribute(self.cache_attributes_bak[a]) + for stage in stream_source.stages: + shape = self._fold_case_state(stage, shape, state) + def _replay_streamed_region( self, stream_source: _PatchStreamSource, diff --git a/konfai/data/sampling.py b/konfai/data/sampling.py index c3f10d43..3d3dd666 100644 --- a/konfai/data/sampling.py +++ b/konfai/data/sampling.py @@ -116,6 +116,18 @@ def _displacement_at(stage: DisplacementStage, world_xyz: torch.Tensor, device: values = torch.tensor(stage.values, dtype=_COORDINATE_DTYPE, device=device) extent_xyz = [int(stage.grid.size_zyx[rank - 1 - axis]) for axis in range(rank)] + if stage.order != 1: + # ITK admits a continuous index within ~4 ulps of the spline's valid-region END by nudging + # it JUST inside (``InsideValidRegion``): the support then fits and the value is the + # continuous limit, the outermost tap's weight vanishing at the boundary. Without the nudge + # the whole last plane of a grid commensurate with the coefficient mesh is silently the + # identity -- and a commensurate grid is exactly what a fitted transform domain produces. + for axis in range(rank): + end = float(extent_xyz[axis] - 2) + ulp = float(np.spacing(np.float64(max(1.0, end)))) + at_end = (index[..., axis] - end).abs() <= 4.0 * ulp + index[..., axis] = index[..., axis].masked_fill(at_end, end - ulp) + base = torch.floor(index) - (stage.order - 1) // 2 taps = stage.order + 1 # The two domains ITK actually implements, and they differ. A BSpline is the identity unless its @@ -411,10 +423,18 @@ def gather( shifted = coordinates_xyz[..., axis] - float(source_starts_zyx[array_axis]) local_axes.append((2.0 * shifted + 1.0) / extent - 1.0) # grid_sample orders the last dimension (x, y, z) -- the mirror of the array axes, which is the - # order the coordinates already arrive in. - sampling = torch.stack([local_axes[rank - 1 - axis] for axis in range(rank)], dim=-1).to(work.dtype) + # order the coordinates already arrive in. The grid counts voxels in float32 WHATEVER the + # payload: a half grid quantizes a coordinate at ~2^-11 of the window extent -- 0.06 voxel on a + # 512 axis, far past the ~1e-5 band above -- and the window is upcast with it because + # grid_sample takes one dtype. sampling_dtype's keep-half trade was measured for the values. + blend_dtype = torch.float32 if work.dtype in (torch.float16, torch.bfloat16) else work.dtype + sampling = torch.stack([local_axes[rank - 1 - axis] for axis in range(rank)], dim=-1).to(blend_dtype) out = torch.nn.functional.grid_sample( - work.unsqueeze(0), sampling.unsqueeze(0), mode="bilinear", padding_mode="border", align_corners=False + work.to(blend_dtype).unsqueeze(0), + sampling.unsqueeze(0), + mode="bilinear", + padding_mode="border", + align_corners=False, ).squeeze(0) return out.masked_fill(~inside.unsqueeze(0), fill).type(source.dtype) diff --git a/konfai/data/transform.py b/konfai/data/transform.py index addc46f7..9ebda1c6 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -580,7 +580,13 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) if isinstance(self.min_value, str): if self.min_value == "min": - min_value = torch.min(tensor_masked) + # Seeded-first, as Normalize reads it: on a streamed path the dispatcher has read + # the CASE's statistic from disk and the tensor in hand is one region of it -- + # computed here, the bound (and what save_clip_min records) would be the region's. + if self.mask is None and "Min" in cache_attribute: + min_value = float(cache_attribute["Min"]) + else: + min_value = torch.min(tensor_masked) elif self.min_value.startswith("percentile:"): try: percentile = float(self.min_value.split(":")[1]) @@ -601,7 +607,10 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) if isinstance(self.max_value, str): if self.max_value == "max": - max_value = torch.max(tensor_masked) + if self.mask is None and "Max" in cache_attribute: + max_value = float(cache_attribute["Max"]) + else: + max_value = torch.max(tensor_masked) elif self.max_value.startswith("percentile:"): try: percentile = float(self.max_value.split(":")[1]) @@ -1192,13 +1201,17 @@ class Resample(TransformInverse): dense solve over the whole grid, and a field solved per region is not the restriction of the field solved once. Store the inverse, or invert it where it is written; - a field with no ``max_displacement`` to size its region from; - - a case that does not meet the target grid anywhere -- the output would be ``fill`` from edge to - edge, and an all-background member is a plausible, wrong contribution to a median. - - Every refusal but the last declares ``WHOLE_VOLUME`` with its reason and the run proceeds on the - whole-volume path, so a chain never breaks over one: it only stops being bounded, and says so in - the plan. A case reaching only PART of the target grid is legal and common -- the rest takes - ``fill`` -- and the plan prints how much of the grid it covers. + - a case that does not meet the target grid anywhere -- judged THROUGH the declared map, so a + stored rigid bridging two scanner frames is not mistaken for disjointness. The output would + be ``fill`` from edge to edge, and an all-background member is a plausible, wrong + contribution to a median. + + A refusal the whole-volume path can serve -- an undeclared field bound, a case with no + geometry -- declares ``WHOLE_VOLUME`` with its reason and the run proceeds assembled: the chain + only stops being bounded, and says so in the plan. One that no route can serve -- a map that + cannot be decoded, read or inverted, or a disjoint case -- refuses as the plan is built, + before a byte is written. A case reaching only PART of the target grid is legal and common -- + the rest takes ``fill`` -- and the plan prints how much of the grid it covers. """ def __init__( @@ -1456,9 +1469,33 @@ def transform_shape(self, group_src: str, name: str, shape: list[int], cache_att self._record(name, [int(extent) for extent in shape], cache_attribute) _source, target = self._target_of(name) if name: + self._require_runnable(name) self._refuse_if_disjoint(name) return [int(extent) for extent in target.size_zyx] + def _require_runnable(self, name: str) -> None: + """Refuse AT PLAN TIME a map neither route can apply. + + A refusal the whole-volume path can serve — an undeclared field bound — stays a locality + answer, and the run proceeds assembled. A stored transform that cannot be decoded, read or + inverted fails the streamed path and the whole-volume one at the same line, so declaring + WHOLE_VOLUME for it would print a plan the run then contradicts by dying per case, after + bytes are written. ``transform_shape`` runs for every case as the plan is built, which is + the earliest the failure is knowable and the only place it costs nothing. + """ + if self.transforms is None: + return + try: + self._stored_stages(name) + except TransformError: + raise + except Exception as error: # a corrupt store fails both routes; name the case and the cure + raise TransformError( + f"'Resample' cannot read the map for case '{name}', so no route can apply it:" + f" {type(error).__name__}: {error}.", + "Check the group names under 'transforms:' and that every case has an entry in each.", + ) from error + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # The geometry is judged on the attribute in hand -- the case's own header, as the base # contract has it -- and not on what the cohort has been seen to carry: one case of a group @@ -1532,8 +1569,10 @@ def stream_region( def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: shape = [int(extent) for extent in tensor.shape[1:]] - if name not in self._grids: - self._record(name, shape, cache_attribute) + # Re-recorded on every call: the whole-volume path is handed the case's own evolved header, + # never a region's, and a grid recorded by an earlier walk may describe the stored volume + # rather than this stage's true input (a Canonical upstream, a Resample before this one). + self._record(name, shape, cache_attribute) source, target = self._grids_of(name) # 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. @@ -1609,26 +1648,51 @@ def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatia def coverage(self, name: str) -> float: """The fraction of the target grid that reads from inside the recorded case.""" - return self._coverage(*self._target_of(name)) + source, target = self._target_of(name) + return self._coverage(source, target, self._map_bound(name)) + + def _map_bound(self, name: str) -> TransformBound | None: + """The declared map's bound, for a coverage judged where the samples actually land. + + ``None`` when there is no map — or when nothing bounds it: a coverage that cannot be judged + must not refuse, and the unboundable configurations carry a fallback reason of their own. + """ + if self.transforms is None and self.displacement is None: + return None + try: + return self._bound(name) + except Exception: # an unreadable or unbounded map answers None, never a crash + return None @classmethod - def _coverage(cls, source: Grid, target: Grid) -> float: + def _coverage(cls, source: Grid, target: Grid, bound: TransformBound | None = None) -> float: """The fraction of ``target`` that reads from inside ``source``, from geometry alone. - The GRID CHANGE only, never the map: a stored transform is what makes the two meet, and - counting the target as uncovered because a registration has not been applied yet would call - every warp disjoint. 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. + 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 + 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. """ axes = [ np.linspace(0.0, float(extent) - 1.0, min(cls._COVERAGE_PROBES, int(extent))) for extent in reversed(target.size_zyx) ] lattice = np.stack([axis.ravel() for axis in np.meshgrid(*axes, indexing="ij")], axis=-1) - index = target.index_to_world.then(source.world_to_index).apply(lattice) + 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) + ) inside = np.ones(index.shape[0], dtype=bool) for axis in range(source.rank): - inside &= (index[:, axis] >= -0.5) & (index[:, axis] < source.size_zyx[source.rank - 1 - axis] - 0.5) + 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]) return float(np.count_nonzero(inside)) / float(inside.size) def _refuse_if_disjoint(self, name: str) -> None: @@ -1663,7 +1727,7 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") if missing & self._target.needs: return None - covered = self._coverage(source, self._target.of(source, name)) + covered = self._coverage(source, self._target.of(source, name), self._map_bound(name)) except TransformError: return None if covered >= self._WORTH_SAYING: @@ -2844,6 +2908,10 @@ def stream_region_source( return source_slices def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + # Nothing to state for a case this cannot reorient: no geometry, or a direction that is not + # 3-D. Its __call__ fails loudly before reaching here; a landing fold must not fail for it. + if not Grid.readable(cache_attribute) or cache_attribute.get_np_array("Direction").size != 9: + return initial_matrix = cache_attribute.get_tensor("Direction").reshape(3, 3).to(torch.double) initial_origin = cache_attribute.get_tensor("Origin") spacing = cache_attribute.get_tensor("Spacing").to(torch.double) diff --git a/tests/unit/test_case_expansion.py b/tests/unit/test_case_expansion.py index e7669e1b..811fe9cd 100644 --- a/tests/unit/test_case_expansion.py +++ b/tests/unit/test_case_expansion.py @@ -456,3 +456,38 @@ def test_a_shared_cache_before_the_marker_is_swept_once_for_every_copy(tmp_path: out = Dataset(tmp_path / "out", "h5") for a in (1, 2, 3): assert out.is_dataset_exist("CT", f"CASE_000_r{a:02d}") + + +def test_interleaved_patch_reads_of_two_copies_each_keep_their_own_grid(tmp_path: Path) -> None: + """Reading copy 1, copy 2, then copy 1 again returns the same bytes every time. + + A stage keys its per-case records by the CASE name — a stored transform is looked up by it — + so the copies of an Expand share one key and the last planned copy's grids would win: with + per-copy Permute draws, ``Resample(shape=[6,6,6])`` folds a DIFFERENT index map per copy, and + a re-read of copy 1 after copy 2's plan silently returned copy 2's sampling of copy 1's data. + """ + from konfai.data.patching import DatasetPatch + from konfai.data.transform import Resample + + source = _source(tmp_path) + permute = _draw(Permute(prob_permute=[0.5, 0.5])) + manager = DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=source, + patch=DatasetPatch(patch_size=[4, 6, 6]), + transforms=[Expand(nb=2, pattern="{name}_r{a:02d}", seed=0), permute, Resample(shape=[6, 6, 6])], + data_augmentations_list=[], + ) + base = torch.from_numpy(source.read_data("CT", "CASE_000")[0].copy()) + + def truth(a: int) -> torch.Tensor: + drawn = permute.compute("CASE_000", 0, a - 1, base.clone()) + return Resample(shape=[6, 6, 6])(f"GT_{a}", drawn, _image_attributes()) + + for index, a in [(0, 1), (0, 2), (0, 1), (1, 2), (1, 1)]: + got = manager.get_data(index, a, [], True) + expected = manager.patch.get_data(truth(a), index, a, True) + assert torch.equal(got, expected), f"patch {index} of copy {a} returned another copy's sampling" diff --git a/tests/unit/test_resample_sampler_rules.py b/tests/unit/test_resample_sampler_rules.py index b36042b2..ee5cecef 100644 --- a/tests/unit/test_resample_sampler_rules.py +++ b/tests/unit/test_resample_sampler_rules.py @@ -190,3 +190,23 @@ def test_a_region_reads_the_same_voxels_as_the_whole_volume() -> None: reach = (slice(None), slice(3, None), slice(4, None), slice(5, None)) span = float(whole.max() - whole.min()) torch.testing.assert_close(partial[reach], whole[reach], rtol=0, atol=1e-5 * span) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the CPU path upcasts half before sampling") +def test_a_cuda_half_volume_is_blended_through_float32_coordinates() -> None: + """The blend's grid counts voxels in float32 whatever the payload. + + A half grid quantizes a coordinate at ~2^-11 of the window extent — 0.06 voxel on a 512 axis, + far past the ~1e-5 streamed-vs-whole band this path claims — so a half CUDA volume must land + each sample where the float32 volume lands it, to half's own value precision. + """ + extent = 512 + # The steepest gradient a volume has: adjacent voxels 1000 apart, so a mis-landed coordinate + # shows as tens of units (measured: 60.0 through a half grid, 0.022 through a float32 one). + x = torch.arange(extent, dtype=torch.float32) + row = torch.where(x % 2 == 0, torch.zeros_like(x), torch.full_like(x, 1000.0)) + volume = row.expand(1, 4, 6, extent).contiguous() + coordinates = _coordinates((4, 6, extent), scales=[0.9, 0.8, 0.97], offsets=[0.1, 0.2, 0.3]).cuda() + exact = gather(volume.cuda(), coordinates, [0, 0, 0], [4, 6, extent], "linear", 0.0) + half = gather(volume.half().cuda(), coordinates, [0, 0, 0], [4, 6, extent], "linear", 0.0) + torch.testing.assert_close(half.float(), exact, rtol=0.0, atol=1.0) diff --git a/tests/unit/test_resample_transform.py b/tests/unit/test_resample_transform.py index 67b3056c..879fa868 100644 --- a/tests/unit/test_resample_transform.py +++ b/tests/unit/test_resample_transform.py @@ -220,22 +220,25 @@ def test_a_case_without_geometry_falls_back_and_says_why(self): assert locality.kind is LocalityKind.WHOLE_VOLUME assert "physical space" in (locality.reason or "") - def test_a_missing_transform_falls_back_and_says_which_group(self): + def test_a_missing_transform_refuses_at_plan_time_and_says_which_group(self): + """A map no route can apply refuses as the plan is built, not per case after bytes. + + Declaring WHOLE_VOLUME instead would print a fallback the run then contradicts by dying: + the whole-volume path needs the same decode this refusal comes from. + """ image = _image() stage = ResampleTransform(transforms={"absent": False}) stage.set_datasets([_StoredTransform("reg", _euler(image))]) - stage.transform_shape("", CASE, list(SIZE), _attribute(image)) - locality = stage.patch_locality(_attribute(image)) - assert locality.kind is LocalityKind.WHOLE_VOLUME - assert "absent" in (locality.reason or "") + with pytest.raises(TransformError, match="absent"): + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) - def test_a_spline_order_with_no_kernel_falls_back_instead_of_crashing_mid_run(self): + def test_a_spline_order_with_no_kernel_refuses_at_plan_time(self): """ITK writes orders 0 and 2 as readily as 3, and neither has a kernel here. - The refusal has to happen where the value is BUILT, not where it is finally sampled: a stage - that decodes such a spline without complaint declares REGRID, passes the plan, and raises on - the first region -- which is halfway through a run, per case, after bytes are already - written. Refused at decode, it is one more whole-volume line in the plan. + The refusal has to happen where the plan is built, not where the value is finally sampled: + a stage that decodes such a spline without complaint passes the plan and raises on the + first region -- halfway through a run, per case, after bytes are already written -- and the + whole-volume path raises the identical error, so there is no fallback to declare. """ image = _image() quadratic = sitk.BSplineTransformInitializer(image, [5] * 3, 2) @@ -244,20 +247,15 @@ def test_a_spline_order_with_no_kernel_falls_back_instead_of_crashing_mid_run(se stage = ResampleTransform(transforms={"reg": False}) stage.set_datasets([_StoredTransform("reg", quadratic)]) - stage.transform_shape("", CASE, list(SIZE), _attribute(image)) - - locality = stage.patch_locality(_attribute(image)) - assert locality.kind is LocalityKind.WHOLE_VOLUME - assert "order 2" in (locality.reason or "") + with pytest.raises(TransformError, match="order 2"): + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) - def test_inverting_a_spline_falls_back_with_the_remedy(self): + def test_inverting_a_spline_refuses_at_plan_time_with_the_remedy(self): image = _image() stage = ResampleTransform(transforms={"reg": True}) stage.set_datasets([_StoredTransform("reg", _bspline(image))]) - stage.transform_shape("", CASE, list(SIZE), _attribute(image)) - locality = stage.patch_locality(_attribute(image)) - assert locality.kind is LocalityKind.WHOLE_VOLUME - assert "Store the inverse" in (locality.reason or "") + with pytest.raises(TransformError, match="Store the inverse"): + stage.transform_shape("", CASE, list(SIZE), _attribute(image)) def test_inverting_a_rigid_map_is_exact_and_still_streams(self): image = _image() @@ -338,3 +336,51 @@ def read_transform(self, group: str, name: str): sitk.Resample(image, image, sitk.CompositeTransform([first, second]), sitk.sitkLinear, 0.0) ) assert float(np.abs(want - got).max()) <= 1e-3 * float(np.abs(want).max()) + + +class _CrossFrameStore: + """One root answering both lookups a cross-frame resample makes: the reference's header and the + stored transform bridging the frames.""" + + def __init__(self, reference: "sitk.Image", transform: "sitk.Transform") -> None: + self.reference = reference + self.transform = transform + + def is_dataset_exist(self, group: str, name: str) -> bool: + del name + return group in ("Reference", "reg") + + def get_infos(self, group: str, name: str): + del group, name + return [1, *list(self.reference.GetSize())[::-1]], _attribute(self.reference) + + def read_transform(self, group: str, name: str) -> "sitk.Transform": + del group, name + return self.transform + + +def test_a_stored_map_bridging_disjoint_frames_is_not_refused_as_disjoint(): + """An MR and a CT can sit 1000 mm apart in stage coordinates with a rigid bridging them. + + The all-fill refusal gates on coverage, and coverage must be judged THROUGH the declared map: + judged before applying it, every cross-frame registration apply — the situation the stage + exists to serve — is refused as disjoint. The counter-assert keeps the gate alive: a map that + leads nowhere still refuses. + """ + from konfai.data.transform import Resample + + 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 + bridge = sitk.TranslationTransform(3, (990.0, 5.0, -2.0)) # target world -> case world + + stage = Resample(reference="ref", reference_group="Reference", transforms={"reg": False}) + stage.set_datasets([_CrossFrameStore(reference, bridge)]) + assert stage.transform_shape("", CASE, list(SIZE), _attribute(case)) == list(SIZE) + assert stage.coverage(CASE) > 0.9 + + astray = sitk.TranslationTransform(3, (500000.0, 0.0, 0.0)) + refused = Resample(reference="ref", reference_group="Reference", transforms={"reg": False}) + refused.set_datasets([_CrossFrameStore(reference, astray)]) + with pytest.raises(TransformError, match="nothing but 'fill'"): + refused.transform_shape("", CASE, list(SIZE), _attribute(case)) diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 4960ba4a..0b395390 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -46,6 +46,7 @@ PatchLocality, Permute, RegionContext, + Resample, ResampleToShape, Softmax, TensorCast, @@ -194,6 +195,108 @@ def test_stream_composed_orientations_with_pointwise_between_match_whole_volume( assert tuple(plans[2].out_shape) == (6, 8) +def _geometry_manager( + stub_class, + volume: np.ndarray, + transforms: list[Transform], + patch: DatasetPatch | None, + spacing: np.ndarray, + direction: np.ndarray, +) -> DatasetManager: + """A manager over the streaming stub with a REAL header — the identity geometry the stub answers + would make every landing fold below trivially right.""" + + class _WithGeometry(stub_class): + def _attributes(self) -> Attribute: + attribute = Attribute() + attribute["Origin"] = np.zeros(volume.ndim - 1) + attribute["Spacing"] = np.asarray(spacing, dtype=np.float64) + attribute["Direction"] = np.asarray(direction, dtype=np.float64).flatten() + return attribute + + return DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, _WithGeometry(volume)), + patch=patch, + transforms=list(transforms), + data_augmentations_list=[], + ) + + +def _fresh_chain_reference(volume: np.ndarray, transforms: list[Transform], attribute: Attribute) -> torch.Tensor: + """The chain run stage by stage on the live header — the semantics every route must reproduce.""" + reference = torch.from_numpy(volume.copy()) + for stage in transforms: + reference = stage("CASE_000", reference, attribute) + return reference + + +def test_a_resample_behind_a_canonical_lands_on_the_reoriented_grid(streaming_dataset_stub) -> None: + """The landing fold evolves the case state, so a Resample is judged on what Canonical left. + + The regression this pins: the fold used to hand every stage the STORED header, so the Resample + recorded the pre-Canonical grid — the whole-volume path then resampled the wrong axis (silently: + every voxel real, the anatomy at the wrong density), and the patched routes crashed or refused + with the blame on the stage. The chain is the shipped TotalSegmentator prediction prefix. + """ + rng = np.random.default_rng(3) + volume = (rng.standard_normal((1, 9, 10, 11)).astype(np.float32)) * 100.0 + # Direction = canonical @ (x<->z swap): Canonical reorients by a signed permutation, after which + # the 2.0 mm axis is x. Resampling to 1.5 iso must therefore widen x: (11, 10, 9*2/1.5=12). + direction = np.diag([-1.0, -1.0, 1.0]) @ np.array([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]]) + spacing = np.asarray([1.5, 1.5, 2.0]) + + def chain() -> list[Transform]: + return [Canonical(), Resample(spacing=[1.5, 1.5, 1.5])] + + whole = _geometry_manager(streaming_dataset_stub, volume, chain(), None, spacing, direction) + assert whole.shapes[0] == [11, 10, 12] + whole.load(whole.transforms, []) + reference = _fresh_chain_reference(volume, chain(), whole.dataset._attributes()) + assert list(reference.shape) == [1, 11, 10, 12] + assert torch.equal(whole.data[0], reference) + + patched = _geometry_manager(streaming_dataset_stub, volume, chain(), DatasetPatch([4, 4, 4]), spacing, direction) + assert patched.can_stream_patch(0), patched.stream_refusal(0) + size = patched.patch.get_size(0) + for index in range(size): + streamed = patched._get_streamed_data(index, 0, True)[0] + # Canonical is an index remap and an axis-aligned spacing change reads one axis at a time on + # global coordinates: both routes are bit-identical to the whole volume, so no tolerance. + assert torch.equal(streamed, patched.patch.get_data(reference, index, 0, True)) + + +def test_a_second_resample_reads_the_first_ones_grid(streaming_dataset_stub) -> None: + """[Resample(3), Resample(1.5)] downsamples then upsamples — the second stage is not a no-op. + + The regression this pins: both stages used to record their grid from the stored header, so the + second saw the ORIGINAL spacing, concluded nothing changes, and handed its input through — the + run then wrote a volume at half the asked density with a header claiming otherwise. + """ + rng = np.random.default_rng(5) + volume = (rng.standard_normal((1, 8, 10, 10)).astype(np.float32)) * 100.0 + + def chain() -> list[Transform]: + return [Resample(spacing=[3.0, 3.0, 3.0]), Resample(spacing=[1.5, 1.5, 1.5])] + + manager = _geometry_manager(streaming_dataset_stub, volume, chain(), None, np.asarray([1.5] * 3), np.eye(3)) + assert manager.shapes[0] == [8, 10, 10] + manager.load(manager.transforms, []) + reference = _fresh_chain_reference(volume, chain(), manager.dataset._attributes()) + assert torch.equal(manager.data[0], reference) + + patched = _geometry_manager( + streaming_dataset_stub, volume, chain(), DatasetPatch([4, 4, 4]), np.asarray([1.5] * 3), np.eye(3) + ) + assert patched.can_stream_patch(0), patched.stream_refusal(0) + for index in range(patched.patch.get_size(0)): + streamed = patched._get_streamed_data(index, 0, True)[0] + assert torch.equal(streamed, patched.patch.get_data(reference, index, 0, True)) + + def test_softmax_channel_axis_is_pointwise_but_spatial_axis_falls_back(build_streaming_manager) -> None: # A channel-axis softmax (dim 0) is spatially pointwise (streamed equality: locality contract). A # softmax over a SPATIAL axis normalises across the whole extent, so a per-patch softmax would @@ -248,6 +351,36 @@ def test_stream_clip_min_max_is_global_stat_and_matches_whole_volume(streaming_d assert stub.full_reads == 0 +def test_a_saved_clip_bound_is_the_cases_statistic_not_the_regions(streaming_dataset_stub) -> None: + """``save_clip_min``/``save_clip_max`` record the bound that was applied — the CASE's. + + On a streamed path the dispatcher seeds the case statistic and the tensor in hand is one + region of it: a bound computed from that region records the region's own extremum on the + attribute, and whatever reads it downstream (an inverse, a Normalize) then works off a number + that depends on which patch happened to run. + """ + rng = np.random.default_rng(4) + volume = (rng.standard_normal((1, 8, 8)).astype(np.float32)) * 100.0 + stage = Clip(min_value="min", max_value="max", save_clip_min=True, save_clip_max=True) + manager = DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, streaming_dataset_stub(volume)), + patch=DatasetPatch([4, 4]), + transforms=[stage], + data_augmentations_list=[], + ) + assert manager.can_stream_patch(0) + # The fixture only pins something if patch 0's extrema differ from the case's. + patch0 = volume[:, :4, :4] + assert float(patch0.min()) != float(volume.min()) or float(patch0.max()) != float(volume.max()) + _tensor, attribute = manager._get_streamed_data(0, 0, True) + assert float(attribute["Min"]) == float(volume.min()) + assert float(attribute["Max"]) == float(volume.max()) + + def test_global_stat_after_float_cast_still_streams_and_matches(build_streaming_manager) -> None: """A value-preserving cast ahead of a GLOBAL_STAT stage must not block streaming. diff --git a/tests/unit/test_transform_bound.py b/tests/unit/test_transform_bound.py index 42da3d3f..a13c4000 100644 --- a/tests/unit/test_transform_bound.py +++ b/tests/unit/test_transform_bound.py @@ -253,3 +253,27 @@ def test_an_affine_stage_carries_no_residual(self): stages = decode_transform_stages(_euler(_image())) assert isinstance(stages[0], AffineStage) np.testing.assert_array_equal(stages[0].bound().residual_xyz, np.zeros(3)) + + def test_the_end_plane_of_a_bsplines_valid_region_is_warped_as_itk_warps_it(self): + """ITK admits a continuous index ON the valid-region end (``InsideValidRegion`` nudges it + back inside), so identity there is wrong bytes. The support slides one control point down, + which changes no value — the outermost tap's weight is exactly zero at an integer offset. + The regression this pins: that plane used to fail the inside test, and a grid commensurate + with its coefficient mesh hits it in whole planes at a time, every voxel silently unmoved. + """ + import torch + from konfai.data.sampling import _displacement_at + + transform = _bspline(_image(oblique=False), mesh=5) + stage = decode_transform_stages(transform)[0] + assert isinstance(stage, DisplacementStage) + extent = np.asarray(list(reversed(stage.grid.size_zyx)), dtype=np.float64) # (x, y, z) + rng = np.random.RandomState(2) + index = rng.uniform(1.0, extent - 2.0 - 1e-6, size=(60, 3)) + for axis in range(3): # each axis in turn pinned exactly on its end plane + index[axis::3, axis] = extent[axis] - 2.0 + world = stage.grid.index_to_world.apply(index) + got = _displacement_at(stage, torch.from_numpy(world), torch.device("cpu")).numpy() + truth = np.array([np.asarray(transform.TransformPoint(tuple(point))) - point for point in world]) + np.testing.assert_allclose(got, truth, rtol=0, atol=1e-12) + assert np.abs(truth).max() > 0.1, "the fixture must displace the end planes, or this pins nothing" From 0bc77759ec575263b6508015237c3d6d9f4b8536 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 23:18:04 +0200 Subject: [PATCH 03/39] feat(transform): route by cost, and a console cut to decisions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming is a memory strategy, not a speed strategy: a halo re-reads its overlap, a regrid pulls each slab's window through its map, and a store without bounded region reads decodes the whole volume once per slab — where loading reads the source once. The plan now prices the streamed route (every pending sweep against its own source, headers only) and a case whose working set fits the per-rank budget is LOADED when streaming would read past 1.5x the source. LOAD is a choice, not a fallback. The predictor's streamed-write threshold prices against the config's resolved per-rank budget instead of instantaneous free RAM, so the same case takes the same route on a loaded machine and an idle one. Dataset.bounded_region_reads is the capability both read. Ten degradations answered WHOLE_VOLUME with no reason; each now names the mechanism and what to change, which is what the plan prints the reason for. Statistics stops being whole-volume: it declares GLOBAL_STAT, reads the seeded case statistic first and computes from the tensor only when nothing seeded one. The console speaks only when the run deviates from the printed plan: a designed refusal prints its message and remedy and exits 1 (traceback under KONFAI_DEBUG=1); the chain line carries the stages and the terminal Write destination; one final line states the counts, the wall time, and where outputs.json is; plan.txt keeps the verbose form. Measured: 23 -> 7 console lines on the 6-output run, log 51 -> 18 lines, the refusal 34 -> 6 lines with the remedy last. --- konfai/data/augmentation.py | 12 +- konfai/data/data_manager.py | 6 +- konfai/data/patching.py | 67 ++++++++ konfai/data/transform.py | 80 ++++++++-- konfai/predictor.py | 26 +++- konfai/transformer.py | 161 +++++++++++++++----- konfai/utils/dataset.py | 49 ++++++ konfai/utils/runtime.py | 18 ++- tests/unit/conftest.py | 4 + tests/unit/test_streamed_read_dispatcher.py | 63 ++++++++ tests/unit/test_transformer_workflow.py | 37 +++++ 11 files changed, 468 insertions(+), 55 deletions(-) diff --git a/konfai/data/augmentation.py b/konfai/data/augmentation.py index c8e1a040..eb2714f9 100755 --- a/konfai/data/augmentation.py +++ b/konfai/data/augmentation.py @@ -514,7 +514,11 @@ def _patch_locality(self, index: int, a: int, cache_attribute: Attribute) -> Pat # what LocalityKind.preserves_statistics lets a later stage trust. Only the draw can say whether # this one is that, and the draw is a property of the copy rather than of the case. if Rotate._index_remap(self.matrix[index][a]) is None: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason="this copy's draw is not a quarter-turn, so it resamples the whole volume;" + " is_quarter: true keeps every draw an index remap", + ) return PatchLocality(LocalityKind.ORIENTATION) def _stream_shape(self, index: int, a: int, shape: list[int]) -> list[int]: @@ -685,7 +689,11 @@ def _patch_locality(self, index: int, a: int, cache_attribute: Attribute) -> Pat # maps values, so a later GLOBAL_STAT could no longer seed from the stored volume -- and only # the tensor's channel count says whether it fires, which a header-time declaration cannot see. if self.vector_field: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason="vector_field: true negates the mirrored component channel, so the stored" + " volume's statistics are not this stage's output's", + ) return PatchLocality(LocalityKind.ORIENTATION) def _stream_region_source( diff --git a/konfai/data/data_manager.py b/konfai/data/data_manager.py index c6977286..818aac2b 100755 --- a/konfai/data/data_manager.py +++ b/konfai/data/data_manager.py @@ -84,7 +84,11 @@ def _format_gib(num_bytes: float) -> str: - return f"{num_bytes / 2**30:.2f} GiB" + """Human bytes at the unit that carries digits — a 0.3 MB refusal must not read '0.00 GiB'.""" + for shift, unit in ((40, "TiB"), (30, "GiB"), (20, "MiB"), (10, "KiB")): + if abs(num_bytes) >= 2**shift: + return f"{num_bytes / 2**shift:.2f} {unit}" + return f"{num_bytes:.0f} B" @dataclass(frozen=True) diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 4fcbd963..6d6c2b00 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -2157,6 +2157,7 @@ def materialize( rewrite: bool = False, fallback_budget_bytes: float | None = None, allow_fallback: bool = True, + prefer_whole: bool = False, ) -> bool: """Write this case's chain to disk by the cheapest path that can, and say which one it took. @@ -2183,6 +2184,14 @@ def materialize( # A chain with an Expand materializes a COPY, whose draw must be part of the plan; without # one it materializes the case itself, and augmentations have nothing to do with writing. apply_augmentations = self._expand is not None and a > 0 + if prefer_whole: + # The plan chose to LOAD: the case fits its budget and streaming would re-read the + # source (:meth:`predicted_stream_read_factor`). A choice, not a fallback -- nothing + # failed -- so ``allow_fallback`` is not consulted; the budget check stays as the belt. + self._enforce_fallback_budget(fallback_budget_bytes) + self._assemble_and_write(a) + self.unload() + return False if self._stream_ready(a, apply_augmentations=apply_augmentations): return True if not allow_fallback: @@ -2243,6 +2252,64 @@ def peak_case_bytes(self) -> int: peak = max(peak, channels * int(np.prod(spatial, dtype=np.int64))) return peak * CASE_ELEMENT_BYTES + def predicted_stream_read_factor(self, a: int = 0, apply_augmentations: bool = False) -> float | None: + """~How many times the streamed route reads the source, priced from the plan alone. + + Streaming is a memory strategy, not a speed strategy: splitting re-reads — a halo re-reads + its overlap, a regrid pulls each slab's window through its map, and a store that cannot + serve bounded region reads decodes the whole volume once per slab — where loading reads the + source once. This is the number the route is CHOSEN with, never the answer: headers only, + one representative slab priced through the plan's own pull maps. ``None`` when the chain + cannot stream (there is no route to price). A chain reading through unmaterialized Save + caches is priced on its top-level plans — a proxy, close enough to route by. + """ + source = self._resolve_patch_stream_source(a, apply_augmentations) + if source is None: + return None + # Each unsatisfied Save sweeps ITS source; past the last boundary the chain reads the + # materialized cache. The route is priced by the dominant segment: a max, not a sum -- + # the segments read different stores, and one that re-reads is the cost either way. + factors = [ + self._segment_read_factor( + sweep.source_dataset, + sweep.source_group, + sweep.source_entry, + [int(extent) for extent in sweep.source_shape], + list(sweep.out_spatial), + sweep.stage_plans, + ) + for sweep in source.pending_sweeps + ] + spatial = [int(extent) for extent in source.shape[1:]] + landed = list(source.stage_plans[-1].out_shape) if source.stage_plans else list(spatial) + factors.append( + self._segment_read_factor( + source.dataset, source.group, source.entry, list(source.shape), landed, source.stage_plans + ) + ) + return max(factors) + + def _segment_read_factor( + self, + dataset: Dataset, + group: str, + entry: str, + source_shape: list[int], + landed: list[int], + plans: tuple[_ReadStagePlan, ...], + ) -> float: + """One segment's reads over its source's voxels, slab by slab through the plan's own pulls.""" + rows = self._sweep_rows(list(landed), int(source_shape[0])) + if not dataset.bounded_region_reads(group, entry): + return float(max(1, -(-landed[0] // rows))) # every slab decodes the whole store + read = 0 + for start in range(0, landed[0], rows): + span = [slice(start, min(start + rows, landed[0])), *(slice(0, extent) for extent in landed[1:])] + for plan in reversed(plans): + span = list(plan.pull(tuple(span))) if plan.pull is not None else span + read += int(np.prod([max(0, part.stop - part.start) for part in span], dtype=np.int64)) + return float(read) / float(max(1, int(np.prod(source_shape[1:], dtype=np.int64)))) + def _enforce_fallback_budget(self, fallback_budget_bytes: float | None) -> None: if fallback_budget_bytes is None: return diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 9ebda1c6..5b9c0246 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -549,14 +549,21 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # both force a whole-volume load. A 'min'/'max' bound needs a global disk statistic # (GLOBAL_STAT); fixed float bounds clip each voxel independently (POINTWISE). if self.mask is not None: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"the bounds are read under mask '{self.mask}', a second whole volume; drop the mask to stream", + ) stat_keys: set[str] = set() for bound, key in ((self.min_value, "Min"), (self.max_value, "Max")): if isinstance(bound, str): if bound.lower() == key.lower(): stat_keys.add(key) else: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"a '{bound}' bound needs the whole histogram; fixed values or" + " 'min'/'max' (a seeded statistic) stream", + ) if not stat_keys: return PatchLocality(LocalityKind.POINTWISE) return PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=frozenset(stat_keys)) @@ -760,7 +767,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # a volume-global disk statistic (GLOBAL_STAT); when both are given, the standardization is a # per-voxel affine map with constant coefficients (POINTWISE). if self.mask is not None: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"the statistics are taken under mask '{self.mask}', a second whole volume;" + " drop the mask to stream", + ) stat_keys: set[str] = set() if self.mean is None: stat_keys.add("Mean") @@ -2031,7 +2042,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # the whole extent, so it falls back to the whole volume. if self.dim == 0: return PatchLocality(LocalityKind.POINTWISE) - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"dim {self.dim} reduces a spatial axis, which spans the whole extent; dim: 0" + " reduces the channels and streams", + ) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: if "number_of_channels_per_model" in cache_attribute: @@ -2135,7 +2150,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # the whole extent, so a per-patch argmax would diverge -- fall back to the whole volume. if self.dim == 0: return PatchLocality(LocalityKind.POINTWISE) - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"dim {self.dim} reduces a spatial axis, which spans the whole extent; dim: 0" + " reduces the channels and streams", + ) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: return torch.argmax(tensor, dim=self.dim).unsqueeze(self.dim) @@ -2151,7 +2170,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # across the whole extent, so a per-patch softmax would diverge -- fall back to the whole volume. if self.dim == 0: return PatchLocality(LocalityKind.POINTWISE) - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=f"dim {self.dim} reduces a spatial axis, which spans the whole extent; dim: 0" + " reduces the channels and streams", + ) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: return torch.softmax(tensor, dim=self.dim) @@ -2879,7 +2902,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # one -- mirroring or permuting -- remaps indices, which is what ORIENTATION streams; an oblique # one is resampled from the whole volume. if self._orthogonal_remap(cache_attribute) is None: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason="the case's direction cosines are oblique (or unreadable), so the" + " reorientation is a resample of the whole volume rather than an index remap", + ) return PatchLocality(LocalityKind.ORIENTATION) def stream_region_source( @@ -2942,7 +2969,11 @@ def _inverse_remap(self, cache_attribute: Attribute) -> list[tuple[int, bool]] | def inverse_patch_locality(self, cache_attribute: Attribute) -> PatchLocality: if self._inverse_remap(cache_attribute) is None: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason="the direction this inverse restores is oblique (or not on the attribute)," + " so the reorientation back is a resample of the whole volume", + ) return PatchLocality(LocalityKind.ORIENTATION) def inverse_transform_shape(self, shape: list[int], cache_attribute: Attribute) -> list[int]: @@ -3395,14 +3426,33 @@ def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) class Statistics(Transform): + """Record the volume's Min/Max/Mean/Std on the case, under ``Image*`` keys. + + Streams: the four numbers are exactly what the disk-statistics scan already computes, so a + streamed chain seeds them (``GLOBAL_STAT``) and each region restates the case's answer instead + of a region's own. + """ + + _KEYS = (("Min", "ImageMin"), ("Max", "ImageMax"), ("Mean", "ImageMean"), ("Std", "ImageStd")) + def __init__(self) -> None: super().__init__() + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + return PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=frozenset({"Min", "Max", "Mean", "Std"})) + def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - cache_attribute["ImageMin"] = tensors.float().min() - cache_attribute["ImageMax"] = tensors.float().max() - cache_attribute["ImageMean"] = tensors.float().mean() - cache_attribute["ImageStd"] = tensors.float().std() + for seeded, recorded in self._KEYS: + if seeded not in cache_attribute: + cache_attribute[recorded] = getattr(tensors.float(), seeded.lower())() + continue + # A seeded statistic arrives as a bare scalar or a one-element array, depending on who + # seeded it; float() reads the first form and get_tensor the second. + raw = cache_attribute[seeded] + try: + cache_attribute[recorded] = float(raw) + except (TypeError, ValueError): + cache_attribute[recorded] = float(cache_attribute.get_tensor(seeded).reshape(-1)[0]) return tensors @@ -3423,7 +3473,11 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # declaration, but a group carries only what its writer stored, and without it there is no # translation to make -- only the read that would find one. if "box" not in cache_attribute: - return PatchLocality(LocalityKind.WHOLE_VOLUME) + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason="the case carries no 'box' yet; the foreground box is computed and recorded" + " as the chain is planned, and only a read can find it", + ) return PatchLocality(LocalityKind.CROP) def stream_region_source( diff --git a/konfai/predictor.py b/konfai/predictor.py index 6ae8f2b1..b53e8f82 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -43,7 +43,13 @@ from konfai import config_file, cuda_visible_devices, konfai_root, predictions_directory from konfai.data.augmentation import DataAugmentation -from konfai.data.data_manager import BatchSample, DataPrediction, DatasetIter +from konfai.data.data_manager import ( + _AUTO_MEMORY_SAFETY_FRACTION, + BatchSample, + DataPrediction, + DatasetIter, + _node_local_ranks, +) from konfai.data.patching import ( Accumulator, PathCombine, @@ -173,6 +179,10 @@ def __init__( self._attributes = dict(entry.split("=", 1) for entry in attributes or []) self.reduction_classpath = reduction self.reduction: Reduction + #: The per-rank budget the streamed-vs-assembled route is priced against — the config's + #: number, pushed by the predictor, so the route is a function of configuration and data + #: rather than of the machine's free memory at the moment a case finalizes. + self._per_rank_budget_bytes: float | None = None self.before_reduction_transforms: list[Transform] = [] self.after_reduction_transforms: list[Transform] = [] @@ -213,6 +223,10 @@ def __init__( self._async_writes = None # decided at the first write, once the device is placed self._writer: _AsyncWriter | None = None + def set_memory_budget(self, budget_bytes: float | None) -> None: + """The per-rank budget the streamed-vs-assembled route is priced against.""" + self._per_rank_budget_bytes = budget_bytes + def _torch_device(self) -> torch.device: """The placed device as ``torch.device`` (``NeedDevice`` may hold a bare CUDA ordinal).""" return torch.device("cuda", self.device) if isinstance(self.device, int) else self.device @@ -710,7 +724,13 @@ def _worth_streaming(self, dataset: DatasetIter, index: int, layer: torch.Tensor stacklevel=2, ) fraction = _STREAM_WORTH_MIN_FRACTION - return assembled >= fraction * available_memory_bytes()[0] + # The config's budget, never the machine's mood: the same case takes the same route on a + # loaded machine and an idle one. A writer no predictor configured (a bare test) prices + # against the auto-budget's own rule. + budget = self._per_rank_budget_bytes + if budget is None: + budget = _AUTO_MEMORY_SAFETY_FRACTION * available_memory_bytes()[0] + return assembled >= fraction * budget def _plan_stream( self, @@ -1906,7 +1926,9 @@ def __init__( self.datasets_filename = [] self.predict_path = predictions_directory() / self.name + per_rank_budget = self.dataset.resolved_budget().per_rank_bytes(_node_local_ranks()) for output_dataset in self.outputs_dataset.values(): + output_dataset.set_memory_budget(per_rank_budget) self.datasets_filename.append(output_dataset.filename) # Rebase under the run directory, re-deriving is_directory: a bare string + "/" would flag an # h5 output as a directory and write the hidden dotfile Predictions//Dataset/.h5. diff --git a/konfai/transformer.py b/konfai/transformer.py index df2fe2f7..64dd1b08 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -29,18 +29,27 @@ import json import os import shutil -from dataclasses import dataclass +import time +from dataclasses import dataclass, field from pathlib import Path from typing import Literal, cast import numpy as np +import torch +import torch.distributed as dist import tqdm from ruamel.yaml import YAML from konfai import config_file, transforms_directory 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 CASE_ELEMENT_BYTES, FALLBACK_INFLIGHT_FACTOR, DatasetManager, save_destination +from konfai.data.patching import ( + _SWEEP_SLAB_ROWS, + CASE_ELEMENT_BYTES, + FALLBACK_INFLIGHT_FACTOR, + DatasetManager, + save_destination, +) from konfai.data.transform import Reduce, Save, Transform, split_expand from konfai.utils.config import apply_config, config from konfai.utils.dataset import Attribute, Dataset @@ -63,7 +72,10 @@ class TransformPlanEntry: case: str # the case name, the copy's entry name behind an Expand, or a reduction's output name group_src: str group_dest: str - verdict: str # "STREAM" | "WHOLE-VOLUME" | "SKIP" | "REDUCE" | "REFUSED" + #: "STREAM" | "LOAD" | "WHOLE-VOLUME" | "SKIP" | "REDUCE" | "REFUSED". LOAD is a choice, not a + #: fallback: the case fits the budget and streaming would re-read the source, so the run + #: assembles it -- same bytes as far as the chain is byte-stable, one read instead of many. + verdict: str reason: str | None case_bytes: int #: The cases folded into one, for a reduction. Empty for an ordinary per-case entry. Printed in @@ -100,6 +112,9 @@ class TransformPlan: #: 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, ...] = () + #: Per (group_src, group_dest): the chain spelled out with its destination — the one fact a + #: reader wants from a plan line ("what runs, and where does it land"). + chain_labels: dict[tuple[str, str], str] = field(default_factory=dict) @property def fallback_entries(self) -> list[TransformPlanEntry]: @@ -121,16 +136,24 @@ def budget_violations(self) -> list[TransformPlanEntry]: estimate is from headers alone and the margin is printed, never hidden: an entry within a few percent of the budget may still exceed it. """ - candidates = [entry for entry in self.entries if entry.verdict in ("WHOLE-VOLUME", "REDUCE")] + candidates = [entry for entry in self.entries if entry.verdict in ("WHOLE-VOLUME", "LOAD", "REDUCE")] return [entry for entry in candidates if entry.working_set_bytes > self.budget_bytes] - def report(self) -> str: - lines = [ + def report(self, verbose: bool = True) -> str: + """The plan as text. ``verbose`` is the plan.txt form; the console gets the same facts with + the estimator caveat only where an estimate actually gates something.""" + estimated = bool(self.fallback_entries or any(entry.reduced for entry in self.entries)) + header = ( f"[Transformer] plan over {self.world_size} rank(s) | per-rank budget" - f" {_format_gib(self.budget_bytes)} ({self.budget_desc}) | fallback working set = case" - f" x {CASE_ELEMENT_BYTES} B x {FALLBACK_INFLIGHT_FACTOR} (in-flight copy), headers-only" - f" estimate | output dtype/channels assumed {self.dtype_hypothesis} until the first slab" - ] + f" {_format_gib(self.budget_bytes)} ({self.budget_desc})" + ) + if verbose or estimated: + header += ( + f" | fallback working set = case x {CASE_ELEMENT_BYTES} B x {FALLBACK_INFLIGHT_FACTOR}" + f" (in-flight copy), headers-only estimate | output dtype/channels assumed" + f" {self.dtype_hypothesis} until the first slab" + ) + lines = [header] for group_src, dropped in sorted(self.dropped_cases.items()): if dropped: lines.append( @@ -142,12 +165,13 @@ def report(self) -> str: for entry in self.entries: by_chain.setdefault((entry.group_src, entry.group_dest), []).append(entry) for (group_src, group_dest), entries in by_chain.items(): + label = self.chain_labels.get((group_src, group_dest), "") + chain = f"{group_src} -> {group_dest}{f' ({label})' if label else ''}" reductions = [entry for entry in entries if entry.reduced] if reductions: for entry in reductions: lines.append( - f" {group_src} -> {group_dest}: REDUCE {len(entry.reduced)} case(s) -> 1" - f" output '{entry.case}' -- {entry.verdict}" + f" {chain}: REDUCE {len(entry.reduced)} case(s) -> 1 output '{entry.case}' -- {entry.verdict}" ) if entry.reason: lines.append(f" {entry.reason}") @@ -158,7 +182,7 @@ def report(self) -> str: ) lines.append(f" cases: {', '.join(entry.reduced)}") continue - counts = dict.fromkeys(("STREAM", "WHOLE-VOLUME", "SKIP"), 0) + counts = dict.fromkeys(("STREAM", "LOAD", "WHOLE-VOLUME", "SKIP"), 0) for entry in entries: counts[entry.verdict] += 1 expanded = [entry for entry in entries if entry.expanded_from] @@ -167,21 +191,22 @@ def report(self) -> str: shared = sum(1 for entry in expanded if entry.regime == "shared") solo = sum(1 for entry in expanded if entry.regime == "solo") lines.append( - f" {group_src} -> {group_dest}: EXPAND {cases} case(s) ->" + f" {chain}: EXPAND {cases} case(s) ->" f" {len(expanded)} cop(ies) -- {shared} STREAM (shared read pass)," f" {solo} STREAM (own pass), {counts['WHOLE-VOLUME']} WHOLE-VOLUME," f" {counts['SKIP']} SKIP (copy already written)" ) else: lines.append( - f" {group_src} -> {group_dest}: {len(entries)} case(s) --" - f" {counts['STREAM']} STREAM, {counts['WHOLE-VOLUME']} WHOLE-VOLUME," + f" {chain}: {len(entries)} case(s) --" + f" {counts['STREAM']} STREAM, {counts['LOAD']} LOAD," + f" {counts['WHOLE-VOLUME']} WHOLE-VOLUME," f" {counts['SKIP']} SKIP (output already written)" ) reasons: dict[str, int] = {} for entry in entries: - if entry.reason and (entry.verdict == "WHOLE-VOLUME" or entry.regime == "solo"): - label = "WHOLE-VOLUME" if entry.verdict == "WHOLE-VOLUME" else "own pass" + if entry.reason and (entry.verdict in ("WHOLE-VOLUME", "LOAD") or entry.regime == "solo"): + label = entry.verdict if entry.verdict in ("WHOLE-VOLUME", "LOAD") else "own pass" reasons[f"{label}: {entry.reason}"] = reasons.get(f"{label}: {entry.reason}", 0) + 1 for reason, count in reasons.items(): lines.append(f" ({count} cop(ies)) {reason}" if expanded else f" ({count} case(s)) {reason}") @@ -225,6 +250,8 @@ def __init__( self._shards: list[list[int]] = [] self._case_names: list[str] = [] self._reductions: dict[str, CaseReduction | None] = {} + self._planned: dict[tuple[str, str], str] = {} + self._sub_cap_sweeps = False def _managers(self) -> dict[str, list[DatasetManager]]: prepared = self.dataset._prepared_data @@ -408,6 +435,33 @@ def output_destinations(self) -> list[dict[str, str]]: ) return destinations + #: Streaming re-reads at most this much of the source before a case that FITS the budget is + #: loaded whole instead. At ~1x the two routes read the same bytes and streaming holds one slab + #: where the load holds the case, so streaming wins the tie; past it the re-reads are the cost. + _STREAM_WORTH_FACTOR = 1.5 + + def _route(self, manager: DatasetManager, case_bytes: int, budget_bytes: float) -> tuple[str, str | None]: + """``STREAM`` or ``LOAD``, priced against the budget — the machine chooses the route, never + the answer. + + A case whose working set exceeds the budget must stream. One that fits streams only while + streaming is no dearer than loading (:meth:`DatasetManager.predicted_stream_read_factor`); + past that it is LOADED — one read of the source instead of many, holding a case the budget + already covers. Expand copies are not routed here: their copies share one read pass, which + amortizes the very re-reads a per-copy load would multiply. + """ + working_set = case_bytes * FALLBACK_INFLIGHT_FACTOR + if working_set <= budget_bytes: + factor = manager.predicted_stream_read_factor(0, apply_augmentations=False) + if factor is not None and factor > self._STREAM_WORTH_FACTOR: + return "LOAD", ( + f"fits the per-rank budget (~{_format_gib(working_set)} vs" + f" {_format_gib(budget_bytes)}); streaming would read ~{factor:.1f}x the source" + ) + if manager._sweep_rows(list(manager.shapes[0]), int(manager.base_shape[0])) < _SWEEP_SLAB_ROWS: + self._sub_cap_sweeps = True + return "STREAM", None + def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> TransformPlan: """Plan every (case, chain) on the launcher, headers plus one write probe per destination.""" budget = self.dataset.resolved_budget() @@ -418,8 +472,14 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor entries: list[TransformPlanEntry] = [] probed: set[tuple[str, str]] = set() planned_dtypes: set[str] = set() + chain_labels: dict[tuple[str, str], str] = {} for group_dest, managers in self._managers().items(): group_src = self._group_src_of(group_dest) + if managers: + names = [type(stage).__name__ for stage in managers[0].transforms] + destination, _terminal_group = self._terminal_destination(managers[0]) + names[-1] += f" {destination.filename}:{destination.file_format}" + chain_labels[(group_src, group_dest)] = " -> ".join(names) reduction = self._reduction(group_dest, managers) if reduction is not None: # Read off the chain, not assumed: a pointwise cast is allowed after the Reduce, and @@ -468,6 +528,8 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor planned_dtypes.add(str(self._dtype_hypothesis(managers[0]))) expansion = split_expand(managers[0].transforms)[1] if managers else None for manager in managers: + # The plan prices the very slabs the run will sweep: same budget, same rows. + manager.set_memory_budget(per_rank_budget) case_bytes = manager.peak_case_bytes() destination, group = self._terminal_destination(manager) if expansion is not None: @@ -506,7 +568,7 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor elif manager.can_stream_patch(0, apply_augmentations=False): probe_failure = self._probe_write_destinations(manager, probed) if probe_failure is None: - verdict, reason = "STREAM", None + verdict, reason = self._route(manager, case_bytes, per_rank_budget) else: # The run would fail the sweep at its first slab and fall back: say so now. verdict, reason = "WHOLE-VOLUME", probe_failure @@ -531,6 +593,7 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor dropped, dtype_hypothesis, tuple(self._plan_notes()), + chain_labels, ) def _plan_notes(self) -> list[str]: @@ -543,6 +606,12 @@ def _plan_notes(self) -> list[str]: case does not: what is printed is the set of distinct things there are to say. """ notes: list[str] = [] + if self._sub_cap_sweeps: + notes.append( + "the budget lowers the sweep slab height below the default; through a" + " non-separable linear resample the streamed values then differ from a" + " taller-slab run by ~1e-5 of the data's range" + ) 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 @@ -581,7 +650,7 @@ def setup(self, world_size: int): self._overwrite = os.environ.get("KONFAI_OVERWRITE", "False") == "True" plan = self.compute_plan(world_size, self._overwrite) report = plan.report() - print(report) + print(plan.report(verbose=False)) # The plan is an artifact, not a log line: Studio filters routine startup lines from the # console stream, and a plan that only ever existed as one is a plan nobody can re-read. (self.transform_path / "plan.txt").write_text(report + "\n") @@ -667,6 +736,9 @@ def setup(self, world_size: int): ) self._budget_bytes = plan.budget_bytes + # The run executes the route the plan priced (LOAD assembles by choice, not fallback), and + # the console only speaks when the run DEVIATES from what the plan already printed. + self._planned = {(entry.group_dest, entry.case): entry.verdict for entry in plan.entries} self._case_names = list(self.dataset._prepared_train_names) # Work items, not cases: an ordinary chain has one per case, a reduction exactly one for all # of them. Each item still writes its own entry, so the disjoint-writes guard above holds. @@ -682,18 +754,21 @@ def setup(self, world_size: int): self.dataloader = [[] for _ in range(max(1, world_size))] def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders): - """Materialize this rank's cases, one at a time, and say which path wrote each.""" + """Materialize this rank's cases. The plan already said what will happen: the console gets + the deviations from it, the live counter, and one final line that says how it went.""" del world_size, local_rank, dataloaders + started = time.monotonic() managers = self._managers() shard = self._shards[global_rank] - counts = {"STREAM": 0, "WHOLE-VOLUME": 0, "SKIP": 0, "REDUCE": 0} + counts = {"STREAM": 0, "LOAD": 0, "WHOLE-VOLUME": 0, "SKIP": 0, "REDUCE": 0} # 'error' holds at run time too: a fallback the plan could not see (a sweep that fails, a # Warp bound exceeded) raises at that case instead of quietly costing a volume. allow_fallback = self.on_fallback != "error" def description() -> str: return ( - f"Transform : {counts['STREAM']} streamed | {counts['REDUCE']} reduced" + f"Transform : {counts['STREAM']} streamed | {counts['LOAD']} loaded" + f" | {counts['REDUCE']} reduced" f" | {counts['WHOLE-VOLUME']} whole-volume | {counts['SKIP']} skipped" ) @@ -708,13 +783,9 @@ def description() -> str: already = reduction.destination.is_dataset_exist(reduction.group, output) if not self._overwrite and already: counts["SKIP"] += 1 - progress.write(f"[Transformer] '{output}' ({group_dest}): SKIP (already written)") else: reduction.materialize(rewrite=self._overwrite) counts["REDUCE"] += 1 - progress.write( - f"[Transformer] '{output}' ({group_dest}): REDUCE {len(reduction.managers)} case(s)" - ) else: manager = group_managers[index] destination, group = self._terminal_destination(manager) @@ -740,26 +811,46 @@ def description() -> str: whole = sum(1 for regime in regimes.values() if regime == "whole-volume") counts["STREAM"] += shared + own counts["WHOLE-VOLUME"] += whole - progress.write( - f"[Transformer] case '{manager.name}' ({group_dest}): EXPAND" - f" {len(copies)} cop(ies) -- {shared} shared pass, {own} own pass," - f" {whole} whole-volume, {len(copies) - len(todo)} skipped" - ) + if whole: + progress.write( + f"[Transformer] case '{manager.name}' ({group_dest}):" + f" {whole} of {len(copies)} cop(ies) took the whole-volume path" + ) elif not self._overwrite and destination.is_dataset_exist(group, manager.name): counts["SKIP"] += 1 - progress.write(f"[Transformer] case '{manager.name}' ({group_dest}): SKIP (already written)") else: + planned = self._planned.get((group_dest, manager.name)) streamed = manager.materialize( 0, rewrite=self._overwrite, fallback_budget_bytes=self._budget_bytes, allow_fallback=allow_fallback, + prefer_whole=planned == "LOAD", ) - verdict = "STREAM" if streamed else "WHOLE-VOLUME" + verdict = "STREAM" if streamed else ("LOAD" if planned == "LOAD" else "WHOLE-VOLUME") counts[verdict] += 1 - progress.write(f"[Transformer] case '{manager.name}' ({group_dest}): {verdict}") + if planned not in (None, verdict): + # The one thing worth a line: the run did NOT do what the plan said. + progress.write( + f"[Transformer] case '{manager.name}' ({group_dest}): {verdict}" + f" (planned {planned}: {manager.stream_refusal(0) or 'see the log'})" + ) progress.set_description(description()) progress.update(1) + totals = dict(counts) + if dist.is_available() and dist.is_initialized(): + gathered = torch.tensor([counts[key] for key in sorted(counts)], dtype=torch.long) + dist.all_reduce(gathered) + totals = {key: int(value) for key, value in zip(sorted(counts), gathered, strict=True)} + if global_rank == 0: + written = totals["STREAM"] + totals["LOAD"] + totals["WHOLE-VOLUME"] + totals["REDUCE"] + resume = f", {totals['SKIP']} already written (--overwrite recomputes)" if totals["SKIP"] else "" + print( + f"[Transformer] done in {time.monotonic() - started:.1f} s: {written} written" + f" ({totals['STREAM']} streamed, {totals['LOAD']} loaded," + f" {totals['WHOLE-VOLUME']} whole-volume, {totals['REDUCE']} reduced){resume}" + f" -> outputs in {self.transform_path / 'outputs.json'}" + ) #: The grammar the strict mode accepts, level by level. ``None`` marks free-form levels (group diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 950190e2..158623e2 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -887,6 +887,17 @@ def file_to_data(self, group: str, name: str) -> tuple[np.ndarray, Attribute]: def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: pass + def bounded_region_reads(self, name: str) -> bool: + """Whether a region read decodes only the region, or the whole volume behind the scenes. + + The base answers ``False``: getting this wrong only ever costs speed, never correctness, + and an unknown backend is priced pessimistically. What it prices is the ROUTE — a store + that decodes the whole volume once per slab makes streaming read the source many times + over, where loading reads it once. + """ + del name + return False + @abstractmethod def file_to_data_statistics( self, @@ -996,6 +1007,10 @@ def file_to_data(self, groups: str, name: str) -> tuple[np.ndarray, Attribute]: dataset.read_direct(data) return data, Attribute(dict(dataset.attrs)) + def bounded_region_reads(self, name: str) -> bool: + del name + return True # h5 is chunked: a slice reads its chunks and nothing else + def file_to_data_slice(self, groups: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: dataset = self._get_dataset(groups, name) data = np.asarray(dataset[slices]) @@ -1328,6 +1343,14 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) - return self._file_to_image_slice(name, path, slices) + def bounded_region_reads(self, name: str) -> bool: + path = self._resolve_data_path(name) + if path is None: + return False + if path.endswith(".npy"): + return True # np.load(mmap) reads the slice off the map + return not path.endswith((".itk.txt", ".fcsv", ".xml", ".vtk")) and self._supports_region_read(path) + def file_to_data_statistics( self, group: str, @@ -1578,6 +1601,10 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) - attributes["Spacing"] = spacing * step_xyz return data, attributes + def bounded_region_reads(self, name: str) -> bool: + del name + return True # zarr is chunked: a slice reads its chunks and nothing else + def file_to_data_statistics( self, group: str, @@ -2108,6 +2135,28 @@ def read_data_slice(self, groups: str, name: str, slices: tuple[slice, ...]) -> raise NameError(f"Dataset entry '{groups}/{name}' not found in {self.filename}.") + def bounded_region_reads(self, groups: str, name: str) -> bool: + """Whether a region read of this entry decodes only the region, or the whole volume. + + What it prices is the ROUTE, never the answer: a store that decodes the whole volume once + per slab (compressed MetaImage, NRRD, gzipped NIfTI) makes streaming read the source many + times over, where loading reads it once. ``False`` for a missing entry — pessimistic, and + only ever costing speed. + """ + if not self._exists_on_disk(): + return False + if self.is_directory: + for sub_directory in self._get_sub_directories(groups): + group = groups.split("/")[-1] + if os.path.exists(f"{self.filename}{sub_directory}{name}{'.h5' if self.file_format == 'h5' else ''}"): + with Dataset.File( + f"{self.filename}{sub_directory}{name}", True, self.file_format, self.level + ) as file: + return file.bounded_region_reads(group) + return False + with Dataset.File(self.filename, True, self.file_format, self.level) as file: + return file.bounded_region_reads(name) + def read_data_statistics( self, groups: str, diff --git a/konfai/utils/runtime.py b/konfai/utils/runtime.py index 437f98d7..213949f8 100644 --- a/konfai/utils/runtime.py +++ b/konfai/utils/runtime.py @@ -61,7 +61,7 @@ statistics_directory, transforms_directory, ) -from konfai.utils.errors import ConfigError +from konfai.utils.errors import ConfigError, KonfAIError from konfai.utils.utils import env_flag @@ -598,6 +598,7 @@ def __init__(self, name: str, rank: int) -> None: # Append, never truncate: this file is opened BEFORE the overwrite prompt runs, so a "w" mode # destroyed the previous run's log even when the user declined the overwrite. self.file = open(self.log_path / f"log_{rank}.txt", "a", buffering=1) + self._last_logged: str | None = None def __enter__(self): super().__enter__() @@ -610,7 +611,12 @@ def __exit__(self, exc_type, exc_val, exc_tb): def write(self, msg: str): super().write(msg) - if self._buffered_line: + # Consecutive identical lines are one fact said twice: a progress bar arrives as several + # write() calls per frame and a case line rides beside its own counter frame, which used to + # multiply the file by ~4x against the console. Only CONSECUTIVE repeats fold -- a fact that + # genuinely recurs later still lands. + if self._buffered_line and self._buffered_line != self._last_logged: + self._last_logged = self._buffered_line self.file.write(self._buffered_line + "\n") self.file.flush() @@ -798,6 +804,14 @@ def wrapper(*args: Any, **kwargs: Any) -> None: ) except KeyboardInterrupt: print("\n[KonfAI] Manual interruption (Ctrl+C)") + except KonfAIError as error: + # A designed refusal: the message says what is wrong and the remedy what to change. + # The traceback under it is framework internals -- 28 lines burying the 3 that matter -- + # so it is shown only to a reader who asked (KONFAI_DEBUG=1). + if env_flag("KONFAI_DEBUG", False): + raise + print(str(error).strip(), file=sys.stderr) + sys.exit(1) finally: if previous_local_ranks is None: os.environ.pop("KONFAI_LOCAL_RANKS", None) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 143d3d03..5734fd16 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -71,6 +71,10 @@ def read_data_statistics(self, group_src: str, name: str, channels: list[int] | "std": float(data.std(ddof=1)), } + def bounded_region_reads(self, group_src: str, name: str) -> bool: + del group_src, name + return True # in-memory: a slice reads the slice + @pytest.fixture def streaming_dataset_stub() -> type[StreamingDatasetStub]: diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 0b395390..3ac16cab 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -381,6 +381,39 @@ def test_a_saved_clip_bound_is_the_cases_statistic_not_the_regions(streaming_dat assert float(attribute["Max"]) == float(volume.max()) +def test_the_predicted_read_factor_prices_the_route_not_the_answer(streaming_dataset_stub) -> None: + """A pointwise chain reads the source once; a store without bounded reads decodes it per slab. + + The factor is what the TRANSFORM verdict routes with: streaming is a memory strategy, and a + case that fits its budget is loaded whole when streaming would read the source many times over. + """ + volume = np.zeros((1, 8, 32, 32), dtype=np.float32) + + def manager(stub) -> DatasetManager: + return DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, stub), + patch=DatasetPatch([4, 8, 8]), + transforms=[Clip(min_value=-10.0, max_value=10.0)], + data_augmentations_list=[], + ) + + bounded = manager(streaming_dataset_stub(volume)) + assert bounded.predicted_stream_read_factor(0) == pytest.approx(1.0) + + class _Unbounded(streaming_dataset_stub): + def bounded_region_reads(self, group_src: str, name: str) -> bool: + return False + + # A tiny budget cuts the sweep into 8 one-row slabs, each decoding the whole store. + unbounded = manager(_Unbounded(volume)) + unbounded.set_memory_budget(1024.0) + assert unbounded.predicted_stream_read_factor(0) == pytest.approx(8.0) + + def test_global_stat_after_float_cast_still_streams_and_matches(build_streaming_manager) -> None: """A value-preserving cast ahead of a GLOBAL_STAT stage must not block streaming. @@ -555,3 +588,33 @@ def test_a_region_stage_recording_geometry_nowhere_the_case_reads_is_refused(bui with pytest.raises(PatchError) as error: manager.get_data(0, 0, [], True) assert "write_stream_cache_attribute" in str(error.value) + + +def test_statistics_streams_off_the_seeded_case_numbers(streaming_dataset_stub) -> None: + """``Statistics`` records the CASE's four numbers — the disk scan already computes them. + + Whole-volume was the default it never needed: seeded, each region restates the case's answer, + and the recorded ``Image*`` keys equal the volume's own statistics rather than a region's. + """ + from konfai.data.transform import Statistics + + rng = np.random.default_rng(6) + volume = (rng.standard_normal((1, 8, 8)).astype(np.float32)) * 100.0 + stub = streaming_dataset_stub(volume) + manager = DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, stub), + patch=DatasetPatch([4, 4]), + transforms=[Statistics()], + data_augmentations_list=[], + ) + assert manager.can_stream_patch(0), manager.stream_refusal(0) + _tensor, attribute = manager._get_streamed_data(0, 0, True) + assert float(attribute["ImageMin"]) == pytest.approx(float(volume.min())) + assert float(attribute["ImageMax"]) == pytest.approx(float(volume.max())) + assert float(attribute["ImageMean"]) == pytest.approx(float(volume.mean()), rel=1e-6) + assert float(attribute["ImageStd"]) == pytest.approx(float(volume.std(ddof=1)), rel=1e-6) + assert stub.full_reads == 0 and stub.stats_reads == 1 diff --git a/tests/unit/test_transformer_workflow.py b/tests/unit/test_transformer_workflow.py index b0411ec5..54b63a09 100644 --- a/tests/unit/test_transformer_workflow.py +++ b/tests/unit/test_transformer_workflow.py @@ -957,3 +957,40 @@ def test_two_ranks_partition_the_cases_and_every_output_is_written_once(tmp_path out = Dataset(f"{tmp_path / 'out_dir'}/", "omezarr") for index in range(3): assert out.is_dataset_exist("CT_out", f"CASE_{index:03d}") + + +def test_a_case_that_fits_is_loaded_when_streaming_would_reread_the_source(tmp_path: Path) -> None: + """The route is chosen from predicted cost against the budget — never the answer. + + A gzipped NIfTI cannot serve bounded region reads, so streaming decodes the whole source once + per slab; a case whose working set fits the budget is then LOADED — one read — and the plan + says so with the factor. LOAD is a choice, not a fallback: on_fallback=error must not refuse + it, and the bytes match the streamed route of the same chain. + """ + rng = np.random.default_rng(3) + source = Dataset(tmp_path / "source", "nii.gz") + volume = (rng.random((1, 8, 32, 32)) * 100).astype(np.float32) + source.write("CT", "CASE_000", volume, _image_attributes()) + out = tmp_path / "out" + transforms = f"""\ + Clip: + min_value: 0.0 + max_value: 50.0 + Write: + dataset: {out}:h5 +""" + config_path = _write_config(tmp_path, transforms, header=" on_fallback: error\n") + # 3x the case: fits (working set = 2x), while the sweep rows drop below the case's extent. + budget = 3 * 8 * 32 * 32 * 4 + config_path.write_text(config_path.read_text().replace("memory_budget: auto", f"memory_budget: {budget}b")) + workflow = _build(tmp_path) + plan = workflow.compute_plan() + entry = next(entry for entry in plan.entries if entry.case == "CASE_000") + assert entry.verdict == "LOAD" + assert "x the source" in (entry.reason or "") + assert not plan.fallback_entries # a choice, not a fallback: on_fallback has nothing to refuse + + workflow.setup(1) + workflow.run_process(1, 0, 0, None) + loaded, _ = Dataset(out, "h5").read_data("CT_out", "CASE_000") + np.testing.assert_array_equal(loaded, np.clip(volume, 0.0, 50.0)) From b970688c3a5db93726bf943acb869cceb671dabe Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 23:18:04 +0200 Subject: [PATCH 04/39] refactor(data): one home per rule; docs, tests and changelog follow - nearest_index, window_index and sampling_dtype move to sampling.py: their only consumers were the two gathers. The dependency arrow is one-way now. - Resample._bound folds its stored stages through bound_of instead of restating the fold inline; composing through the folded affine is interval arithmetic through the product matrix. - Grid.of is from_header plus the refusal, said once. - _DisplacementSource loses its owner/group_keyword parametrisation: Resample is its one owner since the family collapsed. - Dead since the unification, deleted: _halo_from_bound, _array_order_spacing, read_amplification (the route is priced by the plan's own pull maps), and ITK.py's orphaned pre-unification cluster. The docs pages say what the code now does (sampler tolerance band, the LOAD verdict, the env var catalogue). Tests pin the routes that were unpinned: the cost factor monotone in slab fineness, the predictor's stream-worth gate flipping with the config budget, the run handing materialize the plan's route; _sub_cap_sweeps resets per compute_plan. The unreleased changelog section covers the transform release. --- CHANGELOG.md | 65 +++- docs/source/config_guide/transform.md | 12 +- docs/source/reference/environment.md | 6 +- docs/source/usage/large-images.md | 16 +- examples/Transform/README.md | 4 +- examples/Transform/Transform.yml | 2 +- konfai/data/geometry.py | 15 +- konfai/data/sampling.py | 67 ++-- konfai/data/transform.py | 108 +----- konfai/transformer.py | 2 + konfai/utils/ITK.py | 344 ------------------ .../test_konfai_streamed_prediction.py | 10 +- tests/unit/test_itk_transforms.py | 27 -- tests/unit/test_packaging.py | 2 +- tests/unit/test_sampling.py | 27 +- tests/unit/test_streamed_read_dispatcher.py | 27 ++ tests/unit/test_streamed_write_dispatcher.py | 23 ++ tests/unit/test_transformer_workflow.py | 17 +- 18 files changed, 222 insertions(+), 552 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e2bad0b..ba457a19 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,26 @@ written replaces it. - **transform**: `ResampleToShape` needs no geometry at all. A count is a count; only a change of density needs the density it starts from. +- **transform**: the plan chooses the route from predicted cost against the memory budget. + Streaming is a memory strategy — splitting re-reads, loading whole reads once — so a case that + fits the budget is now `LOAD`ed when streaming would re-read the source (a halo re-reads its + overlap, a regrid pulls each slab's window through its map, a compressed store decodes the whole + volume per slab — measured up to 9.7x on an oblique map). The plan prints the predicted factor; + `on_fallback` has nothing to say about a choice. The predictor's streamed-vs-assembled route now + prices the config's budget instead of the machine's free memory at that moment, so the same case + takes the same route on a loaded machine and an idle one. +- **transform**: a console that says what changes a decision. A healthy 6-output run is 7 lines — + the plan's chain lines now carry the stages and the terminal `Write` destination, and one final + line states what was written, how, in how long, and where `outputs.json` is. The run itself only + speaks when it deviates from the printed plan. A designed refusal (`on_fallback: error`, a budget + overrun) prints its message and remedy and exits 1 — 34 lines of framework traceback down to 6; + `KONFAI_DEBUG=1` re-attaches the traceback. Logs stop amplifying progress frames (~4x smaller), + and every byte figure prints at the unit that carries digits instead of `0.00 GiB`. +- **transform**: every configuration-dependent fallback says what to change — a masked `Clip` or + `Standardize`, a percentile bound, a spatial `Sum`/`Argmax`/`Softmax`, an oblique `Canonical`, a + free-angle `Rotate` draw, a vector-field `Flip` — and `Statistics` streams: its four numbers are + the disk scan's own, seeded instead of recomputed per region. + ### 🐛 Fixes - **transform**: a resampled label map no longer comes out shifted against the image beside it. @@ -63,20 +83,55 @@ written replaces it. per array axis from a world displacement, which assumes the direction cosines are the identity; on a turned case the window was short on the axes the displacement actually reached, and a short window returns the border value rather than raising. -- **transform**: a resample refuses what it used to do quietly — resampling in a physical space the - case does not have, applying a transform type nothing bounds, inverting a spline or a field by - building a whole-grid displacement field and iterating on it per case, and warping through a field - with no declared bound. Each declares `WHOLE_VOLUME` with the sentence saying what to change, and - the run proceeds on the whole-volume path. +- **transform**: a resample refuses what it used to do quietly. A refusal the whole-volume path can + serve — an undeclared field bound, a case with no geometry — declares `WHOLE_VOLUME` with the + sentence saying what to change, and the run proceeds assembled. A map neither route can apply — + an unsupported spline order, a missing entry, `invert: true` on a spline or a field — refuses as + the plan is built, before a byte is written; it used to print a fallback the run then contradicted + by dying per case. - **transform**: `ResampleTransform`'s `inverse` defaults to `false`. It always raised `NotImplementedError`, so a prediction finalize through this stage failed at the end of the run rather than at its configuration. +- **data**: a stage is judged on the state the stages before it left, on every landing fold. A + `Resample` behind a `Canonical` recorded the pre-reorientation grid and resampled the wrong axis + — silently, every voxel real, on the exact chain the published TotalSegmentator bundle ships — + and a second `Resample` saw the original spacing and handed its input through as a no-op. Both + now stream, bit-identical to the whole-volume pass. +- **data**: the end plane of a BSpline's valid region is warped as ITK warps it. A grid + commensurate with the coefficient mesh — what a fitted transform domain produces — hit that + plane whole planes at a time, every voxel silently unmoved (2.66 mm of displacement dropped, + against 1e-14 agreement everywhere else). +- **data**: coverage is judged through the declared map before a case is refused as disjoint. An MR + and a CT 1000 mm apart in stage coordinates with a stored rigid bridging them — the situation the + apply step exists to serve — were refused as writing nothing but fill. +- **data**: a half-precision volume on CUDA blends through float32 coordinates. The fused blend + built its sampling grid in the payload's dtype, quantizing a coordinate at ~2^-11 of the window — + 0.06 voxel on a 512 axis, tens of units at a sharp edge (measured 60.0 on a 1000-range fixture, + 0.022 after). +- **data**: interleaved patch reads of two `Expand` copies each keep their own grids; re-reading a + copy after another was planned handed it the other copy's sampling. +- **transform**: `Clip('min'/'max')` clips to the case's seeded statistic, not the region's own — + what `save_clip_min`/`save_clip_max` recorded used to depend on which patch happened to run. + +### ⚡ Performance + +- **data**: a map that factorises is read one axis at a time on global coordinates — most + resamples, bit-identical streamed or whole (CT-sized case: CPU 2269 -> 160 ms, GPU 46.5 -> 2.1 ms, + peak 1.09 -> 0.35 GiB) — and the axis that shrinks most is blended first (9.1x on a thick-slice CT + brought to isotropic). A map that does not factorise (a warp, a rotation, a stored field) goes + through one fused `grid_sample` kernel (4x on a warp); on that path a streamed region and the + whole volume agree to ~1e-5 of the data's range rather than bit for bit, which the plan notes + when a budget shrinks the slabs. + ### 🔧 Internals - **data**: `LocalityKind.RESCALE` is gone. It was the dispatcher's own resample map — a size ratio, which says nothing once a target grid has an origin — and with one resample stage there is one regime, `REGRID`, that the stage owns both halves of. +- **data**: the sampler owns its rules (`nearest_index`, `window_index`, `sampling_dtype` live in + `sampling.py`), `utils/ITK.py` keeps only its live decoders (~340 orphaned pre-unification lines + deleted), and `KONFAI_STREAM_LINEAR_RESAMPLE` — documented, read by nothing — is out of the docs. ## v1.8.0 (2026-08-04) diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index d22da3fe..b56c3876 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -67,8 +67,16 @@ The verdicts, and each one is a fact about *your* run: - **STREAM** — the case is read and written region by region. Memory is one slab, whatever the volume's size. -- **WHOLE-VOLUME** — the case is assembled in memory, then written. Always - correct, never bounded. The line says which stage refused and why. +- **LOAD** — the case *could* stream, fits the budget, and streaming would + re-read the source (a halo re-reads its overlap, a regrid pulls each slab's + window through its map, a compressed store decodes the whole volume per + slab). Loading reads it once; the line prints the predicted factor. A choice, + not a fallback — `on_fallback` has nothing to say about it. Streaming is a + memory strategy: it is chosen when the case does not fit, or when it costs + nothing. +- **WHOLE-VOLUME** — the case cannot stream and is assembled in memory, then + written. Always correct, never bounded. The line says which stage refused and + why. - **SKIP** — the output already exists; nothing is recomputed. - **REDUCE** / **REFUSED** — a chain that folds the cohort into one entry ({doc}`Reduce <../reference/components/transforms>`) diff --git a/docs/source/reference/environment.md b/docs/source/reference/environment.md index 12ca80b6..ff9c8c76 100644 --- a/docs/source/reference/environment.md +++ b/docs/source/reference/environment.md @@ -35,8 +35,7 @@ path or to tune the gate. | Variable | Effect | | --- | --- | | `KONFAI_STREAMED_WRITES` | `0` disables streamed writes entirely (whole-volume reference path). | -| `KONFAI_STREAM_LINEAR_RESAMPLE` | `0` restores bit-exact (non-streamed) linear resample inverses. | -| `KONFAI_STREAM_WORTH_THRESHOLD` | Overrides the "worth streaming" accumulator-size threshold (fraction of allocatable memory). | +| `KONFAI_STREAM_WORTH_THRESHOLD` | Overrides the "worth streaming" accumulator-size threshold (fraction of the per-rank memory budget). Test harnesses set `0` to force the streamed machinery on toy volumes. | | `KONFAI_ASYNC_WRITES` | Controls the background writer for disjoint-file sinks. | | `KONFAI_INLINE_SINGLE_RANK` | Default on. `0` forces a single rank through the spawn path instead of running it in-process — useful when a host process must keep its own CUDA context. | @@ -72,7 +71,8 @@ The codebase also references internal variables such as: - `KONFAI_CONFIG_MODE`, `KONFAI_CONFIG_PATH` — the config binder's mode machine - `KONFAI_APPS_CONFIG` -- `KONFAI_DEBUG`, `KONFAI_DEBUG_LAST_LAYER` +- `KONFAI_DEBUG` — `1` re-attaches the framework traceback to a designed refusal (a + `KonfAIError`), which otherwise prints its message and remedy alone; `KONFAI_DEBUG_LAST_LAYER` - `KONFAI_MASTER_PORT` — distributed rendezvous bookkeeping - `KONFAI_LOCAL_RANKS` — how many ranks share one node's RAM, published by the launcher so a node-scoped `memory_budget` is divided before the spawn. It changes diff --git a/docs/source/usage/large-images.md b/docs/source/usage/large-images.md index ad9b357b..c0beaced 100644 --- a/docs/source/usage/large-images.md +++ b/docs/source/usage/large-images.md @@ -192,14 +192,16 @@ bit-for-bit. `KONFAI_STREAMED_WRITES=0` forces the whole-volume path globally. A **float (linear) resample inverse** — resampling probabilities/logits back to the native grid before an `argmax` — streams within the sliding window by default like the other geometry inverses. On a large multi-class output the -whole-volume `F.interpolate` is otherwise the memory peak (tens of GB), and a +whole-volume resample is otherwise the memory peak (tens of GB), and a `combine: Concat` ensemble makes it worse by keeping every member's channels in -the tensor being resampled; streaming bounds both. It matches the whole-volume -`F.interpolate` to float rounding, not bit-for-bit — `argmax`'d labels absorb it -(a boundary voxel or two may flip; a raw float output differs by ~float-rounding). -Set `KONFAI_STREAM_LINEAR_RESAMPLE=0` to force the exact whole-volume linear -resample when you need bit-identity or the per-member stack; resampling the -`argmax`'d labels (a `nearest`, streaming resample) or collapsing members with +the tensor being resampled; streaming bounds both. An axis-aligned change of +density reads one axis at a time on global coordinates and is **bit-identical** +streamed or whole; only a map that does not factorise (a rotation, a stored +transform, a field) goes through the fused blend, where a streamed region and +the whole volume agree to ~1e-5 of the data's range rather than bit for bit — +`argmax`'d labels absorb it. `KONFAI_STREAMED_WRITES=0` remains the global +whole-volume reference path; resampling the `argmax`'d labels (a `nearest`, +streaming resample, exact by construction) or collapsing members with `combine: Mean`/`Median` first also avoids the peak. ## Verify the behaviour you care about diff --git a/examples/Transform/README.md b/examples/Transform/README.md index ecd90fbf..91a9f24b 100644 --- a/examples/Transform/README.md +++ b/examples/Transform/README.md @@ -41,7 +41,7 @@ not an anatomical frame — and it is why `Reduce` refuses it as stored: case 'CASE_001' lands on extent [44, 60, 52] where 'CASE_000' lands on [48, 56, 56] ``` -`ResampleToReference` is what makes the agreement true rather than waived. It +`Resample: {reference: ...}` is what makes the agreement true rather than waived. It puts every member on one named member's grid — extent, spacing, origin and direction — so `grid: strict` passes because the cohort really is on one grid, not because the check was relaxed. @@ -85,7 +85,7 @@ below it runs once, on the folded result: ```yaml Clip: {min_value: 0.0, max_value: 400.0} # per case -ResampleToReference: {entry: CASE_000, ...} # per case +Resample: {reference: CASE_000, ...} # per case Reduce: {operator: Median, output: template} # <- N becomes 1 here Write: {dataset: ./Template:mha} # once ``` diff --git a/examples/Transform/Transform.yml b/examples/Transform/Transform.yml index e2553c2c..749c9fcb 100644 --- a/examples/Transform/Transform.yml +++ b/examples/Transform/Transform.yml @@ -24,7 +24,7 @@ Transformer: # The cohort as acquired fails `strict`: extents, spacings and origins all differ. # This puts every member on CASE_000's grid, which makes the agreement true rather # than waived. Any member would do -- what matters is that one is named. - ResampleToReference: {entry: CASE_000, group: CT, fill: 0.0} + Resample: {reference: CASE_000, reference_group: CT, fill: 0.0} # The cardinality changes here. Everything above ran once per case; everything below # runs once, on the folded result. Reduce: diff --git a/konfai/data/geometry.py b/konfai/data/geometry.py index ffb0662a..01e6aba3 100644 --- a/konfai/data/geometry.py +++ b/konfai/data/geometry.py @@ -206,24 +206,15 @@ def from_header(cls, spatial_shape: list[int], attribute: Attribute, what: str) @classmethod def of(cls, spatial_shape: list[int], attribute: Attribute, what: str) -> Grid: """The grid a header describes, or a refusal naming ``what`` and the missing key.""" - missing = [key for key in _GEOMETRY_KEYS if key not in attribute] + grid, missing = cls.from_header(spatial_shape, attribute, what) if missing: raise TransformError( - f"The geometry of {what} is needed and its header carries no {', '.join(missing)}.", + f"The geometry of {what} is needed and its header carries no {', '.join(sorted(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).", ) - rank = len(spatial_shape) - origin = _as_float(attribute.get_np_array("Origin"), rank, f"the Origin of {what}", rank) - spacing = _as_float(attribute.get_np_array("Spacing"), rank, f"the Spacing of {what}", rank) - direction = _as_float(attribute.get_np_array("Direction"), rank, f"the Direction of {what}", rank * rank) - if not np.all(spacing > 0.0): - raise TransformError( - f"The Spacing of {what} is {spacing.tolist()}.", - "A spacing is a physical extent per voxel and must be positive on every axis.", - ) - return cls(tuple(int(extent) for extent in spatial_shape), origin, spacing, direction.reshape(rank, rank)) + return grid @staticmethod def readable(attribute: Attribute) -> bool: diff --git a/konfai/data/sampling.py b/konfai/data/sampling.py index 3d3dd666..b92f08ba 100644 --- a/konfai/data/sampling.py +++ b/konfai/data/sampling.py @@ -53,6 +53,49 @@ _COORDINATE_DTYPE = torch.float64 +# -------------------------------------------------------------------------------------------------- +# The rules every sampler below obeys. There are two gather strategies for one arithmetic -- per-axis +# maps where the coordinate is separable, eight flat corners where a displacement makes it not -- and +# the strategies differ for a measured reason. The RULES must not: written out at each site they +# drift, and the drift is silent because a resampled volume looks right either way. + + +def sampling_dtype(tensor: torch.Tensor) -> torch.dtype: + """The dtype to accumulate a weighted sum of ``tensor``'s voxels in. + + An integer input has no arithmetic of its own to interpolate with. A CPU half does, and it should + not be used: torch's CPU Half kernels are missing from older releases and lossy over a sum of + eight terms, at values a scanner actually produces. A CUDA half keeps its own -- every mode has a + Half kernel there, and upcasting a whole multi-class volume would double its memory for nothing. + """ + if not tensor.is_floating_point(): + return torch.float32 + if tensor.device.type == "cpu" and tensor.dtype in (torch.float16, torch.bfloat16): + return torch.float32 + return tensor.dtype + + +def nearest_index(coordinate: torch.Tensor) -> torch.Tensor: + """ITK's nearest: round half UP on the continuous source index. + + ``torch.round`` breaks a tie to the even index, and ``F.interpolate``'s nearest is + ``floor(o * scale)`` -- a statement about a size RATIO, which says nothing once the target grid + carries an origin of its own. On a label map either wrong rule still yields a label map. + """ + return torch.floor(coordinate + 0.5).to(torch.long) + + +def window_index(index: torch.Tensor, n_in: int, region_start: int, window: int) -> torch.Tensor: + """A global source index as an offset into the sub-region that was actually read. + + Clamped twice, and both matter: to the SOURCE first, so a tap past the volume reproduces the + border value rather than wrapping, and to the WINDOW second, so it stays inside the buffer on + hand. The second clamp is only ever load-bearing where the first already put the sample outside, + which the caller masks to fill -- or, for a halo'd read, where the declared bound was checked. + """ + return torch.clamp(torch.clamp(index, 0, n_in - 1) - region_start, 0, window - 1) + + def _kernel_weights(offset: torch.Tensor, order: int) -> torch.Tensor: """The 1-D B-spline weight at distance ``offset`` — ITK's own kernels. @@ -287,8 +330,6 @@ def gather_separable( streamed region against the whole volume, and it is: the per-axis coordinates are global, so a region takes a sub-range of the very numbers the whole volume takes. """ - from konfai.data.transform import nearest_index, sampling_dtype, window_index - rank = len(axes) window_zyx = [int(extent) for extent in source.shape[1:]] extent_zyx = [int(axis.numel()) for axis in axes] @@ -388,8 +429,6 @@ def gather( if not bool(inside.any()): return torch.full(out_shape, fill, device=device, dtype=torch.float32).type(source.dtype) - from konfai.data.transform import nearest_index, sampling_dtype, window_index - work = source.type(sampling_dtype(source)) if mode == "nearest": # One gather, on exact index arithmetic: a nearest pick is discontinuous, so the last bit of @@ -452,23 +491,3 @@ def source_window( grown by the residual), read back as a clamped index window on the source. """ return source_grid.index_window(bound.map_box(target_grid.world_box()), margin) - - -def read_amplification( - target_grid: Grid, - source_grid: Grid, - bound: TransformBound, - regions: list[tuple[slice, ...]], - margin: int = 1, -) -> float: - """How many times the source's voxels this decomposition reads, in total. - - The honest name for what streaming costs: every region's source window is read whole, the - windows overlap, and the finer the decomposition the more they overlap. Monotone in fineness, - so it is a property of the decomposition and not of the transform alone. - """ - total = 0 - for region in regions: - window = source_window(target_grid.sub_grid(region), source_grid, bound, margin) - total += int(np.prod([part.stop - part.start for part in window], dtype=np.int64)) - return float(total) / float(np.prod(source_grid.size_zyx, dtype=np.int64)) diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 5b9c0246..1b724c22 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -43,6 +43,7 @@ Grid, SpatialStages, TransformBound, + bound_of, ) from konfai.data.sampling import ( blend_order, @@ -960,49 +961,6 @@ def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: dict[str, An return tensor.unsqueeze(self.dim) -# -------------------------------------------------------------------------------------------------- -# The rules every sampler below obeys. There are two gather strategies for one arithmetic -- per-axis -# maps where the coordinate is separable, eight flat corners where a displacement makes it not -- and -# the strategies differ for a measured reason. The RULES must not: written out at each site they -# drift, and the drift is silent because a resampled volume looks right either way. - - -def sampling_dtype(tensor: torch.Tensor) -> torch.dtype: - """The dtype to accumulate a weighted sum of ``tensor``'s voxels in. - - An integer input has no arithmetic of its own to interpolate with. A CPU half does, and it should - not be used: torch's CPU Half kernels are missing from older releases and lossy over a sum of - eight terms, at values a scanner actually produces. A CUDA half keeps its own -- every mode has a - Half kernel there, and upcasting a whole multi-class volume would double its memory for nothing. - """ - if not tensor.is_floating_point(): - return torch.float32 - if tensor.device.type == "cpu" and tensor.dtype in (torch.float16, torch.bfloat16): - return torch.float32 - return tensor.dtype - - -def nearest_index(coordinate: torch.Tensor) -> torch.Tensor: - """ITK's nearest: round half UP on the continuous source index. - - ``torch.round`` breaks a tie to the even index, and ``F.interpolate``'s nearest is - ``floor(o * scale)`` -- a statement about a size RATIO, which says nothing once the target grid - carries an origin of its own. On a label map either wrong rule still yields a label map. - """ - return torch.floor(coordinate + 0.5).to(torch.long) - - -def window_index(index: torch.Tensor, n_in: int, region_start: int, window: int) -> torch.Tensor: - """A global source index as an offset into the sub-region that was actually read. - - Clamped twice, and both matter: to the SOURCE first, so a tap past the volume reproduces the - border value rather than wrapping, and to the WINDOW second, so it stays inside the buffer on - hand. The second clamp is only ever load-bearing where the first already put the sample outside, - which the caller masks to fill -- or, for a halo'd read, where the declared bound was checked. - """ - return torch.clamp(torch.clamp(index, 0, n_in - 1) - region_start, 0, window - 1) - - # --------------------------------------------------------------------------------------------- # One resample. Two questions: which grid to write on, and what map to write it through. # --------------------------------------------------------------------------------------------- @@ -1267,9 +1225,7 @@ def __init__( " region a field is read from and means nothing without one.", ) self.displacement: _DisplacementSource | None = ( - _DisplacementSource("Resample", field, field_group, max_displacement, group_keyword="field_group") - if declared - else None + _DisplacementSource(field, field_group, max_displacement) if declared else None ) #: Per case: the grid its own header describes. Recorded where that header is in hand -- #: transform_shape, called for every case as the manager is built. A region read hands back @@ -1469,8 +1425,7 @@ def _bound(self, name: str) -> TransformBound: raise TransformError(self.displacement.undeclared_reason()) folded = TransformBound.shift(np.asarray(declared[:rank], dtype=np.float64)).after(folded) if self.transforms is not None: - for stage in self._stored_stages(name): - folded = stage.bound().after(folded) + folded = bound_of(self._stored_stages(name), rank).after(folded) return folded # ------------------------------------------------------------------ the contract @@ -2270,28 +2225,16 @@ def __init__( class _DisplacementSource: """A displacement field on disk: where it is, how far it reaches, and how to read a region of it. - 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. - - ``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. ``group_keyword`` goes with it: the two owners spell - the field's group differently -- ``Warp`` calls it ``group`` because it has no other, and - ``ResampleToReference`` calls it ``field_group`` because ``group`` is already the reference's. A - remedy naming the wrong one sends the user to change the wrong argument. + :class:`Resample` is its one owner — the family's spellings all resolve to it — so every refusal + speaks as ``Resample`` and names ``field_group``, the argument the user declared. """ def __init__( self, - owner: str, field: str | None, group: str | None, max_displacement: float | str, - group_keyword: str = "group", ) -> None: - self.owner = owner - self.group_keyword = group_keyword # 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. @@ -2301,9 +2244,9 @@ def __init__( self.dataset = Dataset(Path(filename), file_format) elif group is None: raise TransformError( - 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}: {{{group_keyword}: DVF}}.", + "'Resample' has neither a 'field' path nor a group to find the fields in.", + "Name the store — Resample: {field: ./DVF:omezarr} — or, for fields stored beside" + " the cases, the group they are in: Resample: {field_group: DVF}.", ) self.group = group #: The run's own roots, handed over by the owner; only consulted when there is no path. @@ -2314,7 +2257,7 @@ def __init__( max_displacement = float(max_displacement) except ValueError: raise TransformError( - f"'{owner}' has a max_displacement of '{max_displacement}', which is neither a number nor 'auto'.", + f"'Resample' 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 @@ -2380,16 +2323,16 @@ def group_for(self, name: str | None) -> str: 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}: {{{self.group_keyword}: DVF}}.", + "'Resample' has no field store of its own and no group to look for one in.", + "Name the group the fields are in: Resample: {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"'{self.owner}' cannot tell which group of '{self.dataset.filename}' holds {where}: it has {len(groups)}.", - f"Name it: {self.owner}: {{field: ./DVF:omezarr, {self.group_keyword}: DVF}}.", + f"'Resample' cannot tell which group of '{self.dataset.filename}' holds {where}: it has {len(groups)}.", + "Name it: Resample: {field: ./DVF:omezarr, field_group: DVF}.", ) def _root_for(self, name: str | None) -> Dataset: @@ -2401,10 +2344,10 @@ def _root_for(self, name: str | None) -> Dataset: 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"'Resample' 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}}.", + " Give the store a path of its own instead: Resample: {field: ./DVF:omezarr}.", ) def infos(self, name: str) -> tuple[list[int], Attribute]: @@ -2442,25 +2385,13 @@ def check_bound(self, field: torch.Tensor, name: str) -> None: 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.", + f" {component}, beyond the {declared:.3f} 'Resample' 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): @@ -2468,13 +2399,6 @@ def _is_declared_displacement(max_displacement: float | str) -> bool: 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(Resample): """Deprecated spelling of ``Resample: {field: ...}`` — a warp on the case's own grid.""" diff --git a/konfai/transformer.py b/konfai/transformer.py index 64dd1b08..60532db9 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -473,6 +473,8 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor probed: set[tuple[str, str]] = set() planned_dtypes: set[str] = set() chain_labels: dict[tuple[str, str], str] = {} + # A plan is computed from scratch: a note set by an earlier plan of this process must not stick. + self._sub_cap_sweeps = False for group_dest, managers in self._managers().items(): group_src = self._group_src_of(group_dest) if managers: diff --git a/konfai/utils/ITK.py b/konfai/utils/ITK.py index d30517fc..f11c8897 100644 --- a/konfai/utils/ITK.py +++ b/konfai/utils/ITK.py @@ -22,13 +22,11 @@ from typing import TYPE_CHECKING, cast import numpy as np -import torch try: import SimpleITK as sitk except ImportError: sitk = None # type: ignore[assignment] -import torch.nn.functional as F from konfai.utils.errors import TransformError @@ -77,306 +75,6 @@ def read_displacement_field(path: str | Path) -> sitk.Image: return sitk.Cast(field, sitk.sitkVectorFloat64) -def _invert_via_displacement_field(transform: sitk.Transform, image: sitk.Image) -> sitk.DisplacementFieldTransform: - if image is None: - raise TransformError( - "Inverting a non-linear transform requires a reference image to sample the displacement field, " - "but none was provided." - ) - displacement_field_filter = sitk.TransformToDisplacementFieldFilter() - displacement_field_filter.SetReferenceImage(image) - displacement_field = displacement_field_filter.Execute(transform) - iterative_inverse = sitk.IterativeInverseDisplacementFieldImageFilter() - iterative_inverse.SetNumberOfIterations(20) - return sitk.DisplacementFieldTransform(iterative_inverse.Execute(displacement_field)) - - -def _copy_transform(transform_cls: type[sitk.Transform], transform: sitk.Transform, invert: bool) -> sitk.Transform: - transform = transform_cls(transform) - if invert: - transform = transform_cls(transform.GetInverse()) - return transform - - -def _image_like(array: np.ndarray, reference: sitk.Image) -> sitk.Image: - result = sitk.GetImageFromArray(array) - result.CopyInformation(reference) - return result - - -def _open_transform( - transform_files: dict[str | sitk.Transform, bool], image: sitk.Image = None -) -> list[sitk.Transform]: - _require_simpleitk() - transforms: list[sitk.Transform] = [] - - for transform_file, invert in transform_files.items(): - if isinstance(transform_file, str): - transform = sitk.ReadTransform(transform_file + ".itk.txt") - else: - transform = transform_file - if transform.GetName() == "TranslationTransform": - transform = _copy_transform(sitk.TranslationTransform, transform, invert) - elif transform.GetName() == "Euler3DTransform": - transform = _copy_transform(sitk.Euler3DTransform, transform, invert) - elif transform.GetName() == "VersorRigid3DTransform": - transform = _copy_transform(sitk.VersorRigid3DTransform, transform, invert) - elif transform.GetName() == "AffineTransform": - transform = _copy_transform(sitk.AffineTransform, transform, invert) - elif transform.GetName() == "DisplacementFieldTransform": - if invert: - transform = _invert_via_displacement_field(transform, image) - else: - transform = sitk.BSplineTransform(transform) - if invert: - transform = _invert_via_displacement_field(transform, image) - transforms.append(transform) - if len(transforms) == 0: - transforms.append(sitk.Euler3DTransform()) - return transforms - - -def _open_rigid_transform(transform_files: dict[str | sitk.Transform, bool]) -> tuple[np.ndarray, np.ndarray]: - transforms = _open_transform(transform_files) - matrix_result = np.identity(3) - translation_result = np.array([0, 0, 0]) - - for transform in transforms: - if hasattr(transform, "GetMatrix"): - matrix = np.linalg.inv(np.array(transform.GetMatrix(), dtype=np.double).reshape((3, 3))) - translation = -np.asarray(transform.GetTranslation(), dtype=np.double) - center = np.asarray(transform.GetCenter(), dtype=np.double) - else: - matrix = np.eye(len(transform.GetOffset())) - translation = -np.asarray(transform.GetOffset(), dtype=np.double) - center = np.asarray([0] * len(transform.GetOffset()), dtype=np.double) - - translation_center = np.linalg.inv(matrix).dot(matrix.dot(translation - center) + center) - translation_result = np.linalg.inv(matrix_result).dot(translation_center) + translation_result - matrix_result = matrix.dot(matrix_result) - return np.linalg.inv(matrix_result), -translation_result - - -def compose_transform( - transform_files: dict[str | sitk.Transform, bool], image: sitk.Image = None -) -> sitk.CompositeTransform: - transforms = _open_transform(transform_files, image) - result = sitk.CompositeTransform(transforms) - return result - - -def flatten_transform(transform_files: dict[str | sitk.Transform, bool]) -> sitk.AffineTransform: - [matrix, translation] = _open_rigid_transform(transform_files) - transform = sitk.AffineTransform(3) - transform.SetMatrix(matrix.flatten()) - transform.SetTranslation(translation) - return transform - - -def apply_to_image_rigid_transform(image: sitk.Image, transform_files: dict[str | sitk.Transform, bool]) -> sitk.Image: - [matrix, translation] = _open_rigid_transform(transform_files) - matrix = np.linalg.inv(matrix) - translation = -translation - data = sitk.GetArrayFromImage(image) - result = sitk.GetImageFromArray(data) - result.SetDirection(matrix.dot(np.array(image.GetDirection()).reshape((3, 3))).flatten()) - result.SetOrigin(matrix.dot(np.array(image.GetOrigin()) + translation)) - result.SetSpacing(image.GetSpacing()) - return result - - -def apply_to_data_transform(data: np.ndarray, transform_files: dict[str | sitk.Transform, bool]) -> np.ndarray: - transforms = compose_transform(transform_files) - result = np.copy(data) - for i in range(data.shape[0]): - result[i, :] = transforms.TransformPoint(np.asarray(data[i, :], dtype=np.double)) - return result - - -def resample_itk( - image_reference: sitk.Image, - image: sitk.Image, - transform_files: dict[str | sitk.Transform, bool], - mask=False, - default_pixel_value: float | None = None, - torch_resample: bool = False, -) -> sitk.Image: - _require_simpleitk() - if torch_resample: - input_tensor = torch.tensor(sitk.GetArrayFromImage(image)).unsqueeze(0) - vectors = [torch.arange(0, s) for s in input_tensor.shape[1:]] - grids = torch.meshgrid(vectors, indexing="ij") - grid = torch.stack(grids) - grid = torch.unsqueeze(grid, 0) - transform_to_displacement_field_filter = sitk.TransformToDisplacementFieldFilter() - transform_to_displacement_field_filter.SetReferenceImage(image) - transform_to_displacement_field_filter.SetNumberOfThreads(16) - new_locs = grid + torch.tensor( - sitk.GetArrayFromImage( - transform_to_displacement_field_filter.Execute(compose_transform(transform_files, image)) - ) - ).unsqueeze(0).permute(0, 4, 1, 2, 3) - shape = new_locs.shape[2:] - for i in range(len(shape)): - new_locs[:, i, ...] = 2 * (new_locs[:, i, ...] / (shape[i] - 1) - 0.5) - new_locs = new_locs.permute(0, 2, 3, 4, 1) - new_locs = new_locs[..., [2, 1, 0]] - result_data = F.grid_sample( - input_tensor.unsqueeze(0).float(), - new_locs.float(), - align_corners=True, - padding_mode="border", - mode="nearest" if input_tensor.dtype == torch.uint8 else "bilinear", - ).squeeze(0) - result_data = result_data.type(torch.uint8) if input_tensor.dtype == torch.uint8 else result_data - result = sitk.GetImageFromArray(result_data.squeeze(0).numpy()) - result.CopyInformation(image_reference) - return result - else: - return sitk.Resample( - image, - image_reference, - compose_transform(transform_files, image), - sitk.sitkNearestNeighbor if mask else sitk.sitkBSpline, - ( - default_pixel_value - if default_pixel_value is not None - else (0 if mask else int(np.min(sitk.GetArrayFromImage(image)))) - ), - ) - - -def parametermap_to_transform( - path_src: str, -) -> sitk.Transform | list[sitk.Transform]: - _require_simpleitk() - transform = sitk.ReadParameterFile(path_src) - - def array_format(x): - return [float(i) for i in x] - - dimension = int(transform["FixedImageDimension"][0]) - - if transform["Transform"][0] == "EulerTransform": - if dimension == 2: - result = sitk.Euler2DTransform() - else: - result = sitk.Euler3DTransform() - parameters = array_format(transform["TransformParameters"]) - fixed_parameters = [*array_format(transform["CenterOfRotationPoint"]), 0] - elif transform["Transform"][0] == "TranslationTransform": - result = sitk.TranslationTransform(dimension) - parameters = array_format(transform["TransformParameters"]) - fixed_parameters = [] - elif transform["Transform"][0] == "AffineTransform": - result = sitk.AffineTransform(dimension) - parameters = array_format(transform["TransformParameters"]) - fixed_parameters = [*array_format(transform["CenterOfRotationPoint"]), 0] - elif transform["Transform"][0] == "BSplineStackTransform": - parameters = array_format(transform["TransformParameters"]) - grid_size = array_format(transform["GridSize"]) - grid_origin = array_format(transform["GridOrigin"]) - grid_spacing = array_format(transform["GridSpacing"]) - grid_direction = ( - np.asarray(array_format(transform["GridDirection"])).reshape((dimension, dimension)).T.flatten() - ) - fixed_parameters = np.concatenate([grid_size, grid_origin, grid_spacing, grid_direction]) - - nb = int(array_format(transform["Size"])[-1]) - sub = int(np.prod(grid_size)) * dimension - results = [] - for i in range(nb): - result = sitk.BSplineTransform(dimension) - sub_parameters = np.asarray(parameters[i * sub : (i + 1) * sub]) - result.SetFixedParameters(fixed_parameters) - result.SetParameters(sub_parameters) - results.append(result) - return results - elif transform["Transform"][0] == "AffineLogStackTransform": - parameters = array_format(transform["TransformParameters"]) - fixed_parameters = [*array_format(transform["CenterOfRotationPoint"]), 0] - - nb = int(transform["NumberOfSubTransforms"][0]) - sub = dimension * 4 - results = [] - for i in range(nb): - result = sitk.AffineTransform(dimension) - sub_parameters = np.asarray(parameters[i * sub : (i + 1) * sub]) - result.SetFixedParameters(fixed_parameters) - - matrix = torch.from_numpy(sub_parameters[: dimension * dimension].reshape(dimension, dimension)).to( - torch.float64 - ) - matrix_exp = torch.matrix_exp(matrix).cpu().numpy().reshape(-1) - - params = np.concatenate([matrix_exp, sub_parameters[-dimension:]]) - result.SetParameters(params) - results.append(result) - return results - elif transform["Transform"][0] == "BSplineTransform": - result = sitk.BSplineTransform(dimension) - - parameters = array_format(transform["TransformParameters"]) - grid_size = array_format(transform["GridSize"]) - grid_origin = array_format(transform["GridOrigin"]) - grid_spacing = array_format(transform["GridSpacing"]) - grid_direction = np.array(array_format(transform["GridDirection"])).reshape((dimension, dimension)).T.flatten() - fixed_parameters = np.concatenate([grid_size, grid_origin, grid_spacing, grid_direction]) - else: - raise NameError(f"Transform {transform['Transform'][0]} doesn't exist") - result.SetFixedParameters(fixed_parameters) - result.SetParameters(parameters) - return result - - -def _resample(data: torch.Tensor, size: list[int]) -> torch.Tensor: - if data.dtype == torch.uint8: - mode = "nearest" - elif len(data.shape) < 4: - mode = "bilinear" - else: - mode = "trilinear" - return ( - torch.nn.functional.interpolate( - data.type(torch.float32).unsqueeze(0), - size=tuple(reversed(size)), - mode=mode, - ) - .squeeze(0) - .type(data.dtype) - ) - - -def resample_isotropic(image: sitk.Image, spacing: list[float] | None = None) -> sitk.Image: - _require_simpleitk() - spacing = spacing or [1.0, 1.0, 1.0] - resize_factor = [y / x for x, y in zip(spacing, image.GetSpacing(), strict=False)] - result = sitk.GetImageFromArray( - _resample( - torch.tensor(sitk.GetArrayFromImage(image)).unsqueeze(0), - [int(size * factor) for size, factor in zip(image.GetSize(), resize_factor, strict=False)], - ) - .squeeze(0) - .numpy() - ) - result.SetDirection(image.GetDirection()) - result.SetOrigin(image.GetOrigin()) - result.SetSpacing(spacing) - return result - - -def resample_resize(image: sitk.Image, size: list[int] | None = None): - _require_simpleitk() - size = size or [100, 512, 512] - result = sitk.GetImageFromArray( - _resample(torch.tensor(sitk.GetArrayFromImage(image)).unsqueeze(0), size).squeeze(0).numpy() - ) - result.SetDirection(image.GetDirection()) - result.SetOrigin(image.GetOrigin()) - result.SetSpacing([x / y * z for x, y, z in zip(image.GetSize(), size, image.GetSpacing(), strict=False)]) - return result - - def box_with_mask(mask: sitk.Image, label: list[int], dilatations: list[int]) -> np.ndarray: _require_simpleitk() @@ -413,48 +111,6 @@ def crop_with_mask(image: sitk.Image, box: np.ndarray) -> sitk.Image: return result -def format_mask_label(mask: sitk.Image, labels: list[tuple[int, int]]) -> sitk.Image: - _require_simpleitk() - data = sitk.GetArrayFromImage(mask) - result_data = np.zeros_like(data, np.uint8) - - for label_old, label_new in labels: - result_data[np.where(data == label_old)] = label_new - - result = sitk.GetImageFromArray(result_data) - result.CopyInformation(mask) - return result - - -def get_flat_label(mask: sitk.Image, labels: None | list[int] = None) -> sitk.Image: - _require_simpleitk() - data = sitk.GetArrayFromImage(mask) - result_data = np.zeros_like(data, np.uint8) - if labels is not None: - for label in labels: - result_data[np.where(data == label)] = 1 - else: - result_data[np.where(data > 0)] = 1 - result = sitk.GetImageFromArray(result_data) - result.CopyInformation(mask) - return result - - -def clip_and_cast(image: sitk.Image, min_value: float, max_value: float, dtype: np.dtype) -> sitk.Image: - _require_simpleitk() - data = sitk.GetArrayFromImage(image) - data[np.where(data > max_value)] = max_value - data[np.where(data < min_value)] = min_value - result = sitk.GetImageFromArray(data.astype(dtype)) - result.CopyInformation(image) - return result - - -# ------------------------------------------------------------------ decoding a stored transform -# The bridge from a sitk.Transform to konfai.data.geometry's sitk-free stages: everything a -# resample needs to sample and bound the map, as plain numpy, decoded once per case. - - def _linear_map(transform: sitk.Transform) -> AffineMap: """The exact world map of a linear transform: ``T(p) = M p + T(0)``. diff --git a/tests/integration/test_konfai_streamed_prediction.py b/tests/integration/test_konfai_streamed_prediction.py index 3dbf90e8..39b0ce31 100644 --- a/tests/integration/test_konfai_streamed_prediction.py +++ b/tests/integration/test_konfai_streamed_prediction.py @@ -21,8 +21,8 @@ The geometry variants exercise the write dispatcher end to end, one per region kind and then in composition: a ``Canonical`` inverse (ORIENTATION — in-slab mirrors), a ``Padding`` inverse (CROP), a -``ResampleToResolution`` inverse on a uint8 chain (RESCALE, streamed in nearest mode, byte-exact) and on -a float chain (RESCALE, streamed in linear mode, matching the reference to float-rounding), a +``ResampleToResolution`` inverse on a uint8 chain (REGRID, streamed in nearest mode, byte-exact) and on +a float chain (REGRID, streamed in linear mode, matching the reference to float-rounding), a two-inverse pipe, and the full three-inverse stack (crop + rescale + reorient composed, streamed end to end). The TTA variants exercise the slab-synchronized cross-copy reduce: an in-plane flip streams (each copy's window reduced slab by slab), while a slab-axis flip must refuse and complete @@ -117,7 +117,7 @@ def main() -> None: padding: [0, 0, 0, 0, 2, 1] mode: constant inverse: true""", - # RESCALE: the inverse resamples back to the stored grid. + # REGRID: the inverse resamples back to the stored grid. "ResampleLabel": """ transforms: ResampleToResolution: spacing: [0.5, 0.5, -1.0] @@ -136,7 +136,7 @@ def main() -> None: dims: '0' inverse: true""", # The full stack the composition exists for — reorient + resample + pad forward, so the finalize - # chain carries CROP + RESCALE + ORIENTATION in sequence on a uint8 labelmap, streamed end to end. + # chain carries CROP + REGRID + ORIENTATION in sequence on a uint8 labelmap, streamed end to end. "GeometryStack": """ transforms: Canonical: inverse: true @@ -149,7 +149,7 @@ def main() -> None: inverse: true""", } -# ResampleLabel/GeometryStack cast to uint8 before the reduction, so the tensor reaching the RESCALE +# ResampleLabel/GeometryStack cast to uint8 before the reduction, so the tensor reaching the REGRID # stage resamples in nearest mode (byte-exact). ResampleFloat keeps the float chain, so it resamples in # linear mode and matches the reference to float-rounding. _UINT8_BEFORE_REDUCTION = """ before_reduction_transforms: diff --git a/tests/unit/test_itk_transforms.py b/tests/unit/test_itk_transforms.py index 0b6e0ef2..99914818 100644 --- a/tests/unit/test_itk_transforms.py +++ b/tests/unit/test_itk_transforms.py @@ -18,12 +18,9 @@ import numpy as np import pytest -from konfai.utils.errors import TransformError sitk = pytest.importorskip("SimpleITK") -from konfai.utils.ITK import _open_transform, apply_to_data_transform # noqa: E402 - def _identity_displacement_field_transform() -> "sitk.DisplacementFieldTransform": field = sitk.Image(4, 4, 4, sitk.sitkVectorFloat64) @@ -31,30 +28,6 @@ def _identity_displacement_field_transform() -> "sitk.DisplacementFieldTransform return sitk.DisplacementFieldTransform(field) -def test_open_transform_invert_displacement_field_without_image_raises() -> None: - """Inverting a displacement-field transform without a reference image is a typed error, not a crash.""" - transform = _identity_displacement_field_transform() - with pytest.raises(TransformError, match="reference image"): - _open_transform({transform: True}, image=None) - - -def test_open_transform_invert_displacement_field_with_image_succeeds() -> None: - reference = sitk.Image(4, 4, 4, sitk.sitkFloat32) - reference.SetSpacing((1.0, 1.0, 1.0)) - transform = _identity_displacement_field_transform() - result = _open_transform({transform: True}, image=reference) - assert len(result) == 1 - - -def test_apply_to_data_transform_returns_ndarray() -> None: - """apply_to_data_transform returns a numpy array (matching its annotation and callers).""" - points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.double) - translation = sitk.TranslationTransform(3, (10.0, 20.0, 30.0)) - result = apply_to_data_transform(points, {translation: False}) - assert isinstance(result, np.ndarray) - np.testing.assert_allclose(result, points + np.array([10.0, 20.0, 30.0])) - - def test_resample_transform_applies_displacement_in_physical_space() -> None: # ResampleTransform must not add the physical (dx, dy, dz) displacement straight onto a (z, y, x) # voxel-index grid: that transposes x/z and treats millimetres as voxels. A +6 mm translation along diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py index 279ee620..ba28d17f 100644 --- a/tests/unit/test_packaging.py +++ b/tests/unit/test_packaging.py @@ -134,7 +134,7 @@ def test_itk_helper_requires_simpleitk(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(itk_module, "sitk", None) with pytest.raises(TransformError, match="SimpleITK"): - itk_module.resample_resize(None) + itk_module.read_displacement_field("missing.mha") def test_main_module_importable() -> None: diff --git a/tests/unit/test_sampling.py b/tests/unit/test_sampling.py index 8b34fc05..ce76bb3d 100644 --- a/tests/unit/test_sampling.py +++ b/tests/unit/test_sampling.py @@ -26,7 +26,7 @@ import pytest import torch from konfai.data.geometry import Grid, bound_of -from konfai.data.sampling import gather, read_amplification, source_index, source_window +from konfai.data.sampling import gather, source_index, source_window sitk = pytest.importorskip("SimpleITK") @@ -253,28 +253,3 @@ def test_the_window_covers_every_coordinate_the_region_samples(self): assert part.start <= need_low and need_high < part.stop, ( f"{label}: axis {array_axis} needs [{need_low}, {need_high}] outside {part}" ) - - -class TestReadAmplification: - def test_it_grows_as_the_decomposition_gets_finer(self): - # The property the cost model rests on, and the reason a gate is needed at all: streaming - # finer never reads less. Measured on an oblique map, where it runs away fastest. - image = _image() - grid = _grid(image) - bound = bound_of(decode_transform_stages(_affine(image)), 3) - ratios = [] - for rows in (SIZE[0], 8, 3, 1): - regions = [ - (slice(start, min(start + rows, SIZE[0])), slice(0, SIZE[1]), slice(0, SIZE[2])) - for start in range(0, SIZE[0], rows) - ] - ratios.append(read_amplification(grid, grid, bound, regions)) - assert ratios == sorted(ratios), f"amplification must be monotone in fineness, got {ratios}" - assert ratios[-1] > 2.0 * ratios[0] - - def test_a_pure_translation_of_a_whole_volume_reads_about_once(self): - image = _image(oblique=False) - grid = _grid(image) - bound = bound_of(decode_transform_stages(sitk.TranslationTransform(3, (1.0, 1.0, 1.0))), 3) - whole = [(slice(0, SIZE[0]), slice(0, SIZE[1]), slice(0, SIZE[2]))] - assert read_amplification(grid, grid, bound, whole) == pytest.approx(1.0, abs=0.35) diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 3ac16cab..2597dd0b 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -618,3 +618,30 @@ def test_statistics_streams_off_the_seeded_case_numbers(streaming_dataset_stub) assert float(attribute["ImageMean"]) == pytest.approx(float(volume.mean()), rel=1e-6) assert float(attribute["ImageStd"]) == pytest.approx(float(volume.std(ddof=1)), rel=1e-6) assert stub.full_reads == 0 and stub.stats_reads == 1 + + +def test_the_read_factor_grows_as_the_budget_cuts_finer_slabs(streaming_dataset_stub) -> None: + """Streaming finer never reads less — the monotonicity the route's pricing rests on. + + A halo chain re-reads its overlap at every slab boundary, so the factor sits near 1 when one + slab covers the volume and grows as the budget shrinks the slabs. This restates, on the + estimator the verdict actually uses, the property the deleted ``read_amplification`` pinned. + """ + volume = np.zeros((1, 32, 16, 16), dtype=np.float32) + manager = DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, streaming_dataset_stub(volume)), + patch=DatasetPatch([8, 16, 16]), + transforms=[Dilate(2)], + data_augmentations_list=[], + ) + factors = [] + for budget in (None, 64_000.0, 16_000.0, 8_000.0): + manager.set_memory_budget(budget) + factors.append(manager.predicted_stream_read_factor(0)) + assert factors[0] == pytest.approx(1.0, abs=0.2) # one slab: the whole source, once + assert factors == sorted(factors), f"the factor must be monotone in fineness, got {factors}" + assert factors[-1] > 2.0 diff --git a/tests/unit/test_streamed_write_dispatcher.py b/tests/unit/test_streamed_write_dispatcher.py index 9538d83c..83420095 100644 --- a/tests/unit/test_streamed_write_dispatcher.py +++ b/tests/unit/test_streamed_write_dispatcher.py @@ -631,3 +631,26 @@ def test_add_layer_streams_a_full_geometry_stack_through_the_composed_pipe(tmp_p assert streamed.dtype == reference.dtype assert torch.equal(streamed, reference) assert list(streamed.shape) == list(volume.shape) + + +def test_the_stream_worth_gate_prices_the_config_budget_not_the_machine() -> None: + """The streamed-vs-assembled route is a function of configuration and data. + + Same accumulators, two budgets: the verdict must flip with the budget — and with none pushed, + the gate prices the auto-budget's own fraction of the machine, never raw free memory. + """ + from types import SimpleNamespace + + from konfai.predictor import OutSameAsGroupDataset + + writer = OutSameAsGroupDataset.__new__(OutSameAsGroupDataset) + writer.group_dest = "CT" + writer.nb_data_augmentation = 1 + layer = torch.zeros(2, 1, dtype=torch.float32) + dataset = SimpleNamespace(get_dataset_from_index=lambda group, index: SimpleNamespace(shapes=[[64, 64, 64]])) + assembled = 2 * 64 * 64 * 64 * layer.element_size() + + writer.set_memory_budget(assembled) # accumulators are 100% of the budget: worth streaming + assert writer._worth_streaming(dataset, 0, layer) + writer.set_memory_budget(assembled * 1000.0) # a sliver of the budget: assembling costs nothing + assert not writer._worth_streaming(dataset, 0, layer) diff --git a/tests/unit/test_transformer_workflow.py b/tests/unit/test_transformer_workflow.py index 54b63a09..1593418d 100644 --- a/tests/unit/test_transformer_workflow.py +++ b/tests/unit/test_transformer_workflow.py @@ -959,7 +959,9 @@ def test_two_ranks_partition_the_cases_and_every_output_is_written_once(tmp_path assert out.is_dataset_exist("CT_out", f"CASE_{index:03d}") -def test_a_case_that_fits_is_loaded_when_streaming_would_reread_the_source(tmp_path: Path) -> None: +def test_a_case_that_fits_is_loaded_when_streaming_would_reread_the_source( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: """The route is chosen from predicted cost against the budget — never the answer. A gzipped NIfTI cannot serve bounded region reads, so streaming decodes the whole source once @@ -991,6 +993,19 @@ def test_a_case_that_fits_is_loaded_when_streaming_would_reread_the_source(tmp_p assert not plan.fallback_entries # a choice, not a fallback: on_fallback has nothing to refuse workflow.setup(1) + # The bytes of a pointwise chain cannot tell the routes apart, so the ROUTE itself is spied: + # the run must hand materialize the plan's choice, not re-derive its own. + from konfai.data.patching import DatasetManager + + routes: list[bool] = [] + original = DatasetManager.materialize + + def spy(self: DatasetManager, a: int = 0, **kwargs) -> bool: + routes.append(bool(kwargs.get("prefer_whole", False))) + return original(self, a, **kwargs) + + monkeypatch.setattr(DatasetManager, "materialize", spy) workflow.run_process(1, 0, 0, None) + assert routes == [True], "the plan said LOAD and the run must execute it" loaded, _ = Dataset(out, "h5").read_data("CT_out", "CASE_000") np.testing.assert_array_equal(loaded, np.clip(volume, 0.0, 50.0)) From f49a00fea08da704f78f62e53a289825395644fd Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 23:18:04 +0200 Subject: [PATCH 05/39] fix(data): what the pre-release review caught MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three findings, each reproduced by execution before fixing: - The end-of-run totals reduce followed no backend: a CPU tensor under NCCL, so every --gpu TRANSFORM run died on bookkeeping after all its cases were written. The tensor now follows the backend. - apply_to_data_transform (with compose_transform and the _open_transform cluster) is restored to utils/ITK.py: it was deleted as orphaned, and apps/impact_reg — released in lockstep — imports it at module level to warp landmarks. Its tests come back with it. - Clip('min'/'max') and Statistics trusted any 'Min'/'Max'/'Mean'/'Std' key, including the bookkeeping an upstream Normalize pushes for its own inverse: [Normalize, Statistics] recorded the pre-normalization extrema beside a post-normalization mean. The read baselines and the reduction's own seeding pass now carry a chain-scoped StatisticsSeeded marker — set before the persistence snapshot, so it never outlives the read — and the two stages trust a statistic only under it. And a fourth, smaller: an Expand copy's fold state now adopts Crop's box, as the case-level folds do, instead of re-reading the volume once per later fold. --- docs/source/concepts/streaming.md | 7 +- docs/source/reference/environment.md | 4 +- docs/source/usage/large-images.md | 5 +- konfai-mcp/konfai_mcp/catalog.py | 5 ++ konfai/data/case_reduction.py | 1 + konfai/data/geometry.py | 4 +- konfai/data/patching.py | 16 +++- konfai/data/sampling.py | 4 +- konfai/data/transform.py | 34 +++++--- konfai/transformer.py | 9 ++- konfai/utils/ITK.py | 79 ++++++++++++++++++- pixi.lock | 6 +- pyproject.toml | 8 +- tests/unit/test_geometry.py | 15 ++-- tests/unit/test_itk_transforms.py | 27 +++++++ tests/unit/test_streamed_read_dispatcher.py | 4 + tests/unit/test_streamed_write_dispatcher.py | 4 +- .../unit/test_transform_locality_contract.py | 14 +++- .../test_write_pyramid_and_field_bound.py | 2 +- 19 files changed, 201 insertions(+), 47 deletions(-) diff --git a/docs/source/concepts/streaming.md b/docs/source/concepts/streaming.md index 23a275f1..2d0ef30d 100644 --- a/docs/source/concepts/streaming.md +++ b/docs/source/concepts/streaming.md @@ -98,6 +98,7 @@ which region of the file a patch needs. | `HALO` | a bounded neighbourhood of radius `halo` | the patch enlarged by `halo`, cropped after | | `ORIENTATION` | flip or permute | the index-remapped region | | `CROP` | the source region is the target translated | the region — reading it *is* the answer | +| `REGRID` | a change of grid: the source region is the target mapped through it | the mapped region, plus the interpolation taps | | `GLOBAL_STAT` | needs whole-volume `Min`/`Max`/`Mean`/`Std` | the statistic once from disk, then the exact patch | | `SLAB` | a per-voxel value map plus a side effect that needs the slab's place in the volume | nothing on the read path — the dispatcher treats it as `WHOLE_VOLUME`; it streams on the write side (`Mask`, `InferenceStack`) | | `WHOLE_VOLUME` | genuinely needs everything | the volume — the fallback | @@ -113,7 +114,7 @@ is the group's `transforms` followed by the copy's own augmentation draw — one list, so a region transform and a region augmentation compose exactly like two region transforms. -Seven conditions reject streaming: +Six conditions reject streaming: 1. any `WHOLE_VOLUME` **or `SLAB`** declaration — the read dispatcher has no slab context, so it treats `SLAB` as a whole load; @@ -266,10 +267,10 @@ To opt in, override `patch_locality` under three rules: - **Total.** Answer for any case, including one with no metadata. A missing key returns `WHOLE_VOLUME`; it never raises. -`ORIENTATION` and `CROP` must also implement `stream_region_source`, mapping a +`ORIENTATION`, `CROP` and `REGRID` must also implement `stream_region_source`, mapping a target patch to the source region. Declaring a region kind without it raises a `TransformError` on the first patch read. `HALO` needs no remap; the -dispatcher derives their regions. +dispatcher derives its regions. ## Equivalence diff --git a/docs/source/reference/environment.md b/docs/source/reference/environment.md index ff9c8c76..3306f0f7 100644 --- a/docs/source/reference/environment.md +++ b/docs/source/reference/environment.md @@ -72,7 +72,9 @@ The codebase also references internal variables such as: - `KONFAI_CONFIG_MODE`, `KONFAI_CONFIG_PATH` — the config binder's mode machine - `KONFAI_APPS_CONFIG` - `KONFAI_DEBUG` — `1` re-attaches the framework traceback to a designed refusal (a - `KonfAIError`), which otherwise prints its message and remedy alone; `KONFAI_DEBUG_LAST_LAYER` + `KonfAIError`), which otherwise prints its message and remedy alone +- `KONFAI_DEBUG_LAST_LAYER` — set it (empty) before a run and the network appends each module + it enters, so after a crash it names the last layer reached - `KONFAI_MASTER_PORT` — distributed rendezvous bookkeeping - `KONFAI_LOCAL_RANKS` — how many ranks share one node's RAM, published by the launcher so a node-scoped `memory_budget` is divided before the spawn. It changes diff --git a/docs/source/usage/large-images.md b/docs/source/usage/large-images.md index c0beaced..7f93a62a 100644 --- a/docs/source/usage/large-images.md +++ b/docs/source/usage/large-images.md @@ -198,8 +198,9 @@ the tensor being resampled; streaming bounds both. An axis-aligned change of density reads one axis at a time on global coordinates and is **bit-identical** streamed or whole; only a map that does not factorise (a rotation, a stored transform, a field) goes through the fused blend, where a streamed region and -the whole volume agree to ~1e-5 of the data's range rather than bit for bit — -`argmax`'d labels absorb it. `KONFAI_STREAMED_WRITES=0` remains the global +the whole volume agree to ~1e-5 of the data's range rather than bit for bit; +an `argmax` over blended logits usually lands on the same label, though +near-tied logits can flip. `KONFAI_STREAMED_WRITES=0` remains the global whole-volume reference path; resampling the `argmax`'d labels (a `nearest`, streaming resample, exact by construction) or collapsing members with `combine: Mean`/`Median` first also avoids the peak. diff --git a/konfai-mcp/konfai_mcp/catalog.py b/konfai-mcp/konfai_mcp/catalog.py index 91a5e204..655cafca 100644 --- a/konfai-mcp/konfai_mcp/catalog.py +++ b/konfai-mcp/konfai_mcp/catalog.py @@ -29,6 +29,7 @@ import importlib import inspect import os +import types from typing import Any # Kinds backed by "concrete subclasses of a base class defined in a single module". @@ -112,6 +113,10 @@ def _list_subclasses(module_path: str, base_name: str) -> list[dict[str, Any]]: base = getattr(module, base_name) components: list[dict[str, Any]] = [] for name, obj in inspect.getmembers(module, inspect.isclass): + # A subscripted builtin generic (konfai.data.transform.SpatialStages) passes isclass on + # Python 3.10 but is not a class there, and issubclass refuses it. + if isinstance(obj, types.GenericAlias): + continue if obj is base or not issubclass(obj, base): continue if inspect.isabstract(obj) or name.startswith("_"): diff --git a/konfai/data/case_reduction.py b/konfai/data/case_reduction.py index 2f58c462..8ef5674a 100644 --- a/konfai/data/case_reduction.py +++ b/konfai/data/case_reduction.py @@ -157,6 +157,7 @@ def write_into(self, attribute: Attribute) -> None: if self._state is None or not self._state["count"]: raise ReductionError("Statistics were requested over an empty volume.", "Check the output extent.") statistics = _finalize_running_statistics(self._state) + attribute["StatisticsSeeded"] = np.float32(1.0) attribute["Min"] = np.float32(statistics["min"]) attribute["Max"] = np.float32(statistics["max"]) attribute["Mean"] = np.asarray([statistics["mean"]], dtype=np.float32) diff --git a/konfai/data/geometry.py b/konfai/data/geometry.py index 01e6aba3..04af08d4 100644 --- a/konfai/data/geometry.py +++ b/konfai/data/geometry.py @@ -339,8 +339,8 @@ def sub_grid(self, region_zyx: tuple[slice, ...]) -> Grid: The load-bearing line of every streamed regrid: a region left at the volume's origin replays the volume's first slab wherever it lands, and the output still looks like an - image. The origin is ``index_to_world`` of the region's start — the same association of - the same product ITK uses in ``TransformIndexToPhysicalPoint``. + image. The origin is ``index_to_world`` of the region's start — one application of the + parent's own map, never a second association that could land a slab origin apart from it. """ start_xyz = np.array([float(part.start) for part in reversed(region_zyx)]) return Grid( diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 6d6c2b00..d1809900 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -1479,11 +1479,14 @@ def _draw_expand_copies(self, reset_state: bool) -> None: # One draw, every copy at once: state_init IS the per-copy sampler, and it wants the # copies' current grids -- which the stages before it just folded. with _drawn_from(expand.draw_seed, self.index, kind, occurrence): - shapes = stage.state_init(self.index, shapes, attributes) + shapes = stage.state_init(self.index, shapes, foldings) continue for index in range(expand.nb): shapes[index] = self._fold_case_state(stage, shapes[index], foldings[index]) for index in range(expand.nb): + # As at the case-level folds: the box a per-copy Crop computed is a case fact, and + # losing it re-reads the volume once per later fold of that copy. + self._adopt_case_facts(foldings[index], attributes[index]) self.cache_attributes.append(attributes[index]) self.shapes.append(list(shapes[index])) self.patch.load(list(shapes[index]), index + 1) @@ -1746,9 +1749,9 @@ def _plan_stream_region( against — and remaps from — the geometry the stages before it left, and folds the spatial shapes stage by stage; a fold that does not land on ``landing_shape`` (the copy's own grid by default) refuses (the safety net for a stage whose shape map is not declared). Any - ``WHOLE_VOLUME`` declaration, an unreadable ``GLOBAL_STAT``, a ``REGRID`` without a known - ``Spacing`` (or that is not a :class:`Resample`), or a halo too wide to be worth reading - rejects streaming. ``seed_statistics=False`` accepts a missing statistic instead of reading + ``WHOLE_VOLUME`` declaration, an unreadable ``GLOBAL_STAT``, a ``REGRID`` whose own + declaration refuses (a stage answers for its own inputs), or a halo too wide to be worth + reading rejects streaming. ``seed_statistics=False`` accepts a missing statistic instead of reading it — for a chain fed by a cache that is not materialized yet, whose re-resolution seeds it from the real entry. Nothing here names a stage: each declares its own contract, and this is where the declarations are read -- which is why a transform and an augmentation are planned @@ -2800,6 +2803,10 @@ def _get_streamed_data( tensor = torch.from_numpy(data) cache_attribute = Attribute(self.cache_attributes_bak[a]) cache_attribute.update(attributes) + # Says the Min/Max/Mean/Std here are the planner's DISK seeds, not a mid-chain stage's + # own bookkeeping (a Normalize pushes 'Min' for its inverse). Set before keys_before, + # so it never persists past this read. + cache_attribute["StatisticsSeeded"] = 1.0 persist = a not in self._stream_attributes_persisted keys_before = set(cache_attribute.keys()) if persist else set() for stage in stream_source.stages: @@ -2921,6 +2928,7 @@ def _replay_streamed_region( tensor = torch.from_numpy(data) cache_attribute.update(attributes) + cache_attribute["StatisticsSeeded"] = 1.0 # same contract as the pointwise route above keys_before = set(cache_attribute.keys()) for stage, plan, source, target in zip(stages, plans, spans[:-1], spans[1:], strict=True): diff --git a/konfai/data/sampling.py b/konfai/data/sampling.py index b92f08ba..75f821ee 100644 --- a/konfai/data/sampling.py +++ b/konfai/data/sampling.py @@ -429,7 +429,9 @@ def gather( if not bool(inside.any()): return torch.full(out_shape, fill, device=device, dtype=torch.float32).type(source.dtype) - work = source.type(sampling_dtype(source)) + # A nearest pick copies voxels: no blend, no working dtype -- and no float trip for a label, + # whose values above 2**24 a float32 cannot carry back (gather_separable holds the same rule). + work = source if mode == "nearest" else source.type(sampling_dtype(source)) if mode == "nearest": # One gather, on exact index arithmetic: a nearest pick is discontinuous, so the last bit of # a coordinate decides which voxel it lands on. Measured over 300 random grid pairs, a diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 1b724c22..d25bef76 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -523,6 +523,17 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) return result +def _seeded_scalar(cache_attribute: Attribute, key: str) -> float: + """A seeded statistic, as whoever seeded it wrote it: a bare scalar or a one-element array. + + ``float()`` reads the first form and ``get_tensor`` the second. + """ + try: + return float(cache_attribute[key]) + except (TypeError, ValueError): + return float(cache_attribute.get_tensor(key).reshape(-1)[0]) + + class Clip(Transform): """Clip tensor intensities to a fixed or data-dependent value range.""" @@ -557,7 +568,7 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: stat_keys: set[str] = set() for bound, key in ((self.min_value, "Min"), (self.max_value, "Max")): if isinstance(bound, str): - if bound.lower() == key.lower(): + if bound == key.lower(): # exactly as __call__ matches it; "MIN" is refused there stat_keys.add(key) else: return PatchLocality( @@ -591,8 +602,8 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) # Seeded-first, as Normalize reads it: on a streamed path the dispatcher has read # the CASE's statistic from disk and the tensor in hand is one region of it -- # computed here, the bound (and what save_clip_min records) would be the region's. - if self.mask is None and "Min" in cache_attribute: - min_value = float(cache_attribute["Min"]) + if self.mask is None and "StatisticsSeeded" in cache_attribute and "Min" in cache_attribute: + min_value = _seeded_scalar(cache_attribute, "Min") else: min_value = torch.min(tensor_masked) elif self.min_value.startswith("percentile:"): @@ -615,8 +626,8 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) if isinstance(self.max_value, str): if self.max_value == "max": - if self.mask is None and "Max" in cache_attribute: - max_value = float(cache_attribute["Max"]) + if self.mask is None and "StatisticsSeeded" in cache_attribute and "Max" in cache_attribute: + max_value = _seeded_scalar(cache_attribute, "Max") else: max_value = torch.max(tensor_masked) elif self.max_value.startswith("percentile:"): @@ -1000,6 +1011,8 @@ class _DerivedGrid(_TargetGrid): """The case's own grid at another density — a spacing, or a count, and where it sits.""" def __init__(self, spacing: list[float] | None, shape: list[int] | None, align: str) -> None: + # A value <= 0 is the KEEP-THIS-AXIS sentinel, normalised to 0 here: that axis takes the + # source's own density/extent in `of`, so a request rescales only the axes it names. self.spacing = None if spacing is None else np.asarray([max(0.0, float(value)) for value in spacing]) self.shape = None if shape is None else tuple(max(0, int(value)) for value in shape) self.align = align @@ -3366,17 +3379,12 @@ def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: return PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=frozenset({"Min", "Max", "Mean", "Std"})) def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: + trusted = "StatisticsSeeded" in cache_attribute for seeded, recorded in self._KEYS: - if seeded not in cache_attribute: + if not trusted or seeded not in cache_attribute: cache_attribute[recorded] = getattr(tensors.float(), seeded.lower())() continue - # A seeded statistic arrives as a bare scalar or a one-element array, depending on who - # seeded it; float() reads the first form and get_tensor the second. - raw = cache_attribute[seeded] - try: - cache_attribute[recorded] = float(raw) - except (TypeError, ValueError): - cache_attribute[recorded] = float(cache_attribute.get_tensor(seeded).reshape(-1)[0]) + cache_attribute[recorded] = _seeded_scalar(cache_attribute, seeded) return tensors diff --git a/konfai/transformer.py b/konfai/transformer.py index 60532db9..408264cc 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -142,7 +142,9 @@ def budget_violations(self) -> list[TransformPlanEntry]: def report(self, verbose: bool = True) -> str: """The plan as text. ``verbose`` is the plan.txt form; the console gets the same facts with the estimator caveat only where an estimate actually gates something.""" - estimated = bool(self.fallback_entries or any(entry.reduced for entry in self.entries)) + estimated = bool( + self.fallback_entries or any(entry.reduced or entry.verdict == "LOAD" for entry in self.entries) + ) header = ( f"[Transformer] plan over {self.world_size} rank(s) | per-rank budget" f" {_format_gib(self.budget_bytes)} ({self.budget_desc})" @@ -841,7 +843,10 @@ def description() -> str: progress.update(1) totals = dict(counts) if dist.is_available() and dist.is_initialized(): - gathered = torch.tensor([counts[key] for key in sorted(counts)], dtype=torch.long) + # NCCL reduces only device tensors; gloo takes CPU ones. Follow the backend, or every + # --gpu run dies on this bookkeeping reduce after all its cases were written. + device = torch.device("cuda") if dist.get_backend() == "nccl" else torch.device("cpu") + gathered = torch.tensor([counts[key] for key in sorted(counts)], dtype=torch.long, device=device) dist.all_reduce(gathered) totals = {key: int(value) for key, value in zip(sorted(counts), gathered, strict=True)} if global_rank == 0: diff --git a/konfai/utils/ITK.py b/konfai/utils/ITK.py index f11c8897..2115b435 100644 --- a/konfai/utils/ITK.py +++ b/konfai/utils/ITK.py @@ -75,6 +75,81 @@ def read_displacement_field(path: str | Path) -> sitk.Image: return sitk.Cast(field, sitk.sitkVectorFloat64) +def _invert_via_displacement_field(transform: sitk.Transform, image: sitk.Image) -> sitk.DisplacementFieldTransform: + if image is None: + raise TransformError( + "Inverting a non-linear transform requires a reference image to sample the displacement field, " + "but none was provided." + ) + displacement_field_filter = sitk.TransformToDisplacementFieldFilter() + displacement_field_filter.SetReferenceImage(image) + displacement_field = displacement_field_filter.Execute(transform) + iterative_inverse = sitk.IterativeInverseDisplacementFieldImageFilter() + iterative_inverse.SetNumberOfIterations(20) + return sitk.DisplacementFieldTransform(iterative_inverse.Execute(displacement_field)) + + +def _copy_transform(transform_cls: type[sitk.Transform], transform: sitk.Transform, invert: bool) -> sitk.Transform: + transform = transform_cls(transform) + if invert: + transform = transform_cls(transform.GetInverse()) + return transform + + +def _image_like(array: np.ndarray, reference: sitk.Image) -> sitk.Image: + result = sitk.GetImageFromArray(array) + result.CopyInformation(reference) + return result + + +def _open_transform( + transform_files: dict[str | sitk.Transform, bool], image: sitk.Image = None +) -> list[sitk.Transform]: + _require_simpleitk() + transforms: list[sitk.Transform] = [] + + for transform_file, invert in transform_files.items(): + if isinstance(transform_file, str): + transform = sitk.ReadTransform(transform_file + ".itk.txt") + else: + transform = transform_file + if transform.GetName() == "TranslationTransform": + transform = _copy_transform(sitk.TranslationTransform, transform, invert) + elif transform.GetName() == "Euler3DTransform": + transform = _copy_transform(sitk.Euler3DTransform, transform, invert) + elif transform.GetName() == "VersorRigid3DTransform": + transform = _copy_transform(sitk.VersorRigid3DTransform, transform, invert) + elif transform.GetName() == "AffineTransform": + transform = _copy_transform(sitk.AffineTransform, transform, invert) + elif transform.GetName() == "DisplacementFieldTransform": + if invert: + transform = _invert_via_displacement_field(transform, image) + else: + transform = sitk.BSplineTransform(transform) + if invert: + transform = _invert_via_displacement_field(transform, image) + transforms.append(transform) + if len(transforms) == 0: + transforms.append(sitk.Euler3DTransform()) + return transforms + + +def compose_transform( + transform_files: dict[str | sitk.Transform, bool], image: sitk.Image = None +) -> sitk.CompositeTransform: + transforms = _open_transform(transform_files, image) + result = sitk.CompositeTransform(transforms) + return result + + +def apply_to_data_transform(data: np.ndarray, transform_files: dict[str | sitk.Transform, bool]) -> np.ndarray: + transforms = compose_transform(transform_files) + result = np.copy(data) + for i in range(data.shape[0]): + result[i, :] = transforms.TransformPoint(np.asarray(data[i, :], dtype=np.double)) + return result + + def box_with_mask(mask: sitk.Image, label: list[int], dilatations: list[int]) -> np.ndarray: _require_simpleitk() @@ -176,7 +251,9 @@ def decode_transform_stages(transform: sitk.Transform) -> SpatialStages: if isinstance(transform, sitk.CompositeTransform): stages: list[AffineStage | DisplacementStage] = [] for index in reversed(range(transform.GetNumberOfTransforms())): - stages.extend(decode_transform_stages(transform.GetNthTransform(index))) + # Downcast restores the member's concrete type where GetNthTransform hands back the + # generic wrapper, which the isinstance dispatch below cannot read. + stages.extend(decode_transform_stages(transform.GetNthTransform(index).Downcast())) return tuple(stages) from konfai.data.geometry import AffineStage diff --git a/pixi.lock b/pixi.lock index 7be40634..01c9c96d 100644 --- a/pixi.lock +++ b/pixi.lock @@ -1985,11 +1985,11 @@ packages: - lxml - requests - huggingface-hub - - simpleitk ; extra == 'itk' + - simpleitk>=2.0 ; extra == 'itk' - h5py ; extra == 'hdf5' - nvidia-ml-py ; extra == 'monitoring' - tensorboard ; extra == 'tensorboard' - - simpleitk ; extra == 'imaging' + - simpleitk>=2.0 ; extra == 'imaging' - h5py ; extra == 'imaging' - pydicom ; extra == 'imaging' - zarr ; extra == 'imaging' @@ -2022,7 +2022,7 @@ packages: - konfai[omezarr] ; extra == 'all' - konfai[export] ; extra == 'all' - konfai[smp] ; extra == 'all' - - simpleitk ; extra == 'dev' + - simpleitk>=2.0 ; extra == 'dev' - h5py ; extra == 'dev' - pydicom ; extra == 'dev' - zarr ; extra == 'dev' diff --git a/pyproject.toml b/pyproject.toml index 1b261dd2..7f053293 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,11 +52,11 @@ konfai = "konfai.main:main" konfai-cluster = "konfai.main:cluster" [project.optional-dependencies] -itk = ["SimpleITK"] +itk = ["SimpleITK>=2.0"] hdf5 = ["h5py"] monitoring = ["nvidia-ml-py"] tensorboard = ["tensorboard"] -imaging = ["SimpleITK", "h5py", "pydicom", "zarr", "ngff-zarr>=0.38", "dask"] +imaging = ["SimpleITK>=2.0", "h5py", "pydicom", "zarr", "ngff-zarr>=0.38", "dask"] vtk = ["vtk"] lpips = ["lpips"] smp = ["segmentation-models-pytorch"] @@ -82,7 +82,7 @@ all = [ "konfai[smp]" ] dev = [ - "SimpleITK", + "SimpleITK>=2.0", "h5py", "pydicom", "zarr", @@ -159,7 +159,7 @@ twine = "*" mypy = "*" types-requests = "*" submitit = "*" -SimpleITK = "*" +SimpleITK = ">=2.0" h5py = "*" tensorboard = "*" nvidia-ml-py = "*" diff --git a/tests/unit/test_geometry.py b/tests/unit/test_geometry.py index 999c59a6..8ce8d7b3 100644 --- a/tests/unit/test_geometry.py +++ b/tests/unit/test_geometry.py @@ -65,16 +65,21 @@ def test_index_to_world_is_transform_index_to_physical_point(self): got = grid.index_to_world.apply(np.asarray(index_xyz)) np.testing.assert_allclose(got, want, rtol=0.0, atol=1e-12) - def test_sub_grid_origin_is_the_slab_origin_bit_for_bit(self): - # The load-bearing line: same product, same association as ITK, on an oblique grid with - # anisotropic spacing and a non-round origin. + def test_sub_grid_origin_is_the_slab_origin(self): + # The load-bearing identity is internal, and bit-for-bit: the slab origin IS the parent + # map applied to the slab's first index -- a recomputation that associated differently + # would shear a streamed regrid at its seams. Agreement with ITK itself holds only to + # tolerance: numpy does not reproduce ITK's product bit-for-bit on every ISA (arm64 + # contracts with FMA where x86-64 does not), so that half shares the neighbouring + # test's 1e-12. grid = _grid(_oblique_direction()) image = _image(grid) region = (slice(11, 21), slice(0, 33), slice(5, 48)) - want = np.array(image.TransformIndexToPhysicalPoint([5, 0, 11])) sub = grid.sub_grid(region) assert sub.size_zyx == (10, 33, 43) - np.testing.assert_array_equal(sub.origin_xyz, want) + np.testing.assert_array_equal(sub.origin_xyz, grid.index_to_world.apply(np.array([5.0, 0.0, 11.0]))) + want = np.array(image.TransformIndexToPhysicalPoint([5, 0, 11])) + np.testing.assert_allclose(sub.origin_xyz, want, rtol=0.0, atol=1e-12) def test_world_to_index_round_trips(self): grid = _grid(_oblique_direction()) diff --git a/tests/unit/test_itk_transforms.py b/tests/unit/test_itk_transforms.py index 99914818..0b6e0ef2 100644 --- a/tests/unit/test_itk_transforms.py +++ b/tests/unit/test_itk_transforms.py @@ -18,9 +18,12 @@ import numpy as np import pytest +from konfai.utils.errors import TransformError sitk = pytest.importorskip("SimpleITK") +from konfai.utils.ITK import _open_transform, apply_to_data_transform # noqa: E402 + def _identity_displacement_field_transform() -> "sitk.DisplacementFieldTransform": field = sitk.Image(4, 4, 4, sitk.sitkVectorFloat64) @@ -28,6 +31,30 @@ def _identity_displacement_field_transform() -> "sitk.DisplacementFieldTransform return sitk.DisplacementFieldTransform(field) +def test_open_transform_invert_displacement_field_without_image_raises() -> None: + """Inverting a displacement-field transform without a reference image is a typed error, not a crash.""" + transform = _identity_displacement_field_transform() + with pytest.raises(TransformError, match="reference image"): + _open_transform({transform: True}, image=None) + + +def test_open_transform_invert_displacement_field_with_image_succeeds() -> None: + reference = sitk.Image(4, 4, 4, sitk.sitkFloat32) + reference.SetSpacing((1.0, 1.0, 1.0)) + transform = _identity_displacement_field_transform() + result = _open_transform({transform: True}, image=reference) + assert len(result) == 1 + + +def test_apply_to_data_transform_returns_ndarray() -> None: + """apply_to_data_transform returns a numpy array (matching its annotation and callers).""" + points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.double) + translation = sitk.TranslationTransform(3, (10.0, 20.0, 30.0)) + result = apply_to_data_transform(points, {translation: False}) + assert isinstance(result, np.ndarray) + np.testing.assert_allclose(result, points + np.array([10.0, 20.0, 30.0])) + + def test_resample_transform_applies_displacement_in_physical_space() -> None: # ResampleTransform must not add the physical (dx, dy, dz) displacement straight onto a (z, y, x) # voxel-index grid: that transposes x/z and treats millimetres as voxels. A +6 mm translation along diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 2597dd0b..3c7d46bc 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -33,6 +33,7 @@ import numpy as np import pytest import torch +from konfai.data import patching from konfai.data.augmentation import DataAugmentationsList from konfai.data.augmentation import Flip as FlipAugmentation from konfai.data.patching import DatasetManager, DatasetPatch @@ -642,6 +643,9 @@ def test_the_read_factor_grows_as_the_budget_cuts_finer_slabs(streaming_dataset_ for budget in (None, 64_000.0, 16_000.0, 8_000.0): manager.set_memory_budget(budget) factors.append(manager.predicted_stream_read_factor(0)) + # The ~1.0 first factor holds only while the no-budget sweep covers the volume in ONE slab; + # a smaller default cap would split it and re-read the Dilate halo at each boundary. + assert patching._SWEEP_SLAB_ROWS >= volume.shape[1], "the first factor's premise moved" assert factors[0] == pytest.approx(1.0, abs=0.2) # one slab: the whole source, once assert factors == sorted(factors), f"the factor must be monotone in fineness, got {factors}" assert factors[-1] > 2.0 diff --git a/tests/unit/test_streamed_write_dispatcher.py b/tests/unit/test_streamed_write_dispatcher.py index 83420095..e93a01b5 100644 --- a/tests/unit/test_streamed_write_dispatcher.py +++ b/tests/unit/test_streamed_write_dispatcher.py @@ -633,7 +633,7 @@ def test_add_layer_streams_a_full_geometry_stack_through_the_composed_pipe(tmp_p assert list(streamed.shape) == list(volume.shape) -def test_the_stream_worth_gate_prices_the_config_budget_not_the_machine() -> None: +def test_the_stream_worth_gate_prices_the_config_budget_not_the_machine(monkeypatch) -> None: """The streamed-vs-assembled route is a function of configuration and data. Same accumulators, two budgets: the verdict must flip with the budget — and with none pushed, @@ -641,6 +641,8 @@ def test_the_stream_worth_gate_prices_the_config_budget_not_the_machine() -> Non """ from types import SimpleNamespace + monkeypatch.delenv("KONFAI_STREAM_WORTH_THRESHOLD", raising=False) + from konfai.predictor import OutSameAsGroupDataset writer = OutSameAsGroupDataset.__new__(OutSameAsGroupDataset) diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index 252bb59b..c5f5a798 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -34,11 +34,11 @@ """ import inspect +import types from dataclasses import dataclass import numpy as np import pytest -import SimpleITK as sitk import torch from konfai.data import augmentation as augmentation_module from konfai.data import transform as transform_module @@ -82,7 +82,7 @@ from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import TransformError -pytest.importorskip("SimpleITK") +sitk = pytest.importorskip("SimpleITK") _CASE_NAME = "CASE_000" _SPATIAL = (9, 10, 11) @@ -325,8 +325,13 @@ def _builtin_transforms() -> list[type[Transform]]: """Every concrete transform class KonfAI ships.""" return [ cls + # A subscripted builtin generic (SpatialStages) passes isclass on Python 3.10 but is not a + # class there, and issubclass refuses it. for _, cls in inspect.getmembers(transform_module, inspect.isclass) - if issubclass(cls, Transform) and cls.__module__ == transform_module.__name__ and not inspect.isabstract(cls) + if not isinstance(cls, types.GenericAlias) + and issubclass(cls, Transform) + and cls.__module__ == transform_module.__name__ + and not inspect.isabstract(cls) ] @@ -598,7 +603,8 @@ def _builtin_augmentations() -> list[type[DataAugmentation]]: return [ cls for _, cls in inspect.getmembers(augmentation_module, inspect.isclass) - if issubclass(cls, DataAugmentation) + if not isinstance(cls, types.GenericAlias) + and issubclass(cls, DataAugmentation) and cls.__module__ == augmentation_module.__name__ and not inspect.isabstract(cls) ] diff --git a/tests/unit/test_write_pyramid_and_field_bound.py b/tests/unit/test_write_pyramid_and_field_bound.py index 7dc91cab..481761fe 100644 --- a/tests/unit/test_write_pyramid_and_field_bound.py +++ b/tests/unit/test_write_pyramid_and_field_bound.py @@ -190,7 +190,6 @@ def test_a_field_records_its_own_bound_on_both_write_paths(tmp_path): assert DISPLACEMENT_BOUND_ATTRIBUTE in Dataset(tmp_path / "streamed", "omezarr").get_infos("DVF", "case")[1] -@_needs_rfc5 def tmp_field_store(field: np.ndarray) -> Path: import tempfile @@ -201,6 +200,7 @@ def tmp_field_store(field: np.ndarray) -> Path: return root / "fields" +@_needs_rfc5 def test_the_recorded_bound_reaches_each_axis_by_its_own_spacing_under_anisotropy() -> None: """What the bound is FOR, read by the stage that consumes it. From 5907865f181d6899a5431b66b5b89dfa9e4b9d16 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 00:03:54 +0200 Subject: [PATCH 06/39] feat(impact-reg): give the CLI its working directory, as the other apps have register/eval/uncertainty staged under tempfile.gettempdir() with no way to say otherwise, so a caller whose TMPDIR is a tmpfs paid volume-sized intermediates in RAM and had no recourse but to override TMPDIR for the whole subprocess -- which also moves what TMPDIR legitimately owns (torch's DataLoader worker sockets, whose AF_UNIX addresses cap at ~108 bytes). --tmp-dir is what every other KonfAI app CLI already exposes through build_app_cli; impact-reg is the only orchestrator that builds the konfai-apps command line by hand (one subprocess per preset, because konfai keeps process-global state) and that hand- built list had dropped this one flag of the ten it forwards. Forwarding it to konfai-apps also removes a full-size write: given a caller-owned workspace, konfai-apps writes straight into -o instead of staging ./Predictions and copying it in (see _stage_result_dir / _collect_result). The moved image and the displacement field are now written once per preset rather than twice. Default unchanged: with no --tmp-dir the commands stage where they always did. --- CHANGELOG.md | 4 ++ apps/impact_reg/impact_reg_konfai/cli.py | 24 +++++++++++ .../impact_reg_konfai/impact_reg.py | 40 +++++++++++++++++-- 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba457a19..3e06c56a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -148,6 +148,10 @@ written replaces it. - **data**: Vote, the reduction operator that folds segmentations without inventing a label - **data**: declare an OME-Zarr pyramid from a Write, and let a field carry its own bound - **impact-reg**: seed the rigid from the centre of mass, not only the frame +- **impact-reg**: `--tmp-dir` on register/eval/uncertainty, the option the other app CLIs already carry — + a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk + instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the + displacement field once per run instead of twice - **studio**: bundle icons through the app interface, and a way to stop Studio (#75) - **examples**: a Transform example -- a template folded out of a cohort, and drawn copies of a case diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index b75dbf9a..23e618a0 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -44,6 +44,24 @@ def _default_preset() -> str: return presets[0] +def _add_tmp_dir(parser: argparse.ArgumentParser) -> None: + """Add ``--tmp-dir``, the option every other KonfAI app CLI exposes through ``build_app_cli``. + + Left unset the command stages under the system temporary directory, as before. It is worth naming + when that directory is the wrong medium: what is staged is volume-sized (the moved image and the + displacement field, written before being collected into ``--output``), so a tmpfs TMPDIR charges it + to RAM. Pointing this beside ``--output`` also puts the intermediates on the results' own filesystem. + """ + parser.add_argument( + "--tmp-dir", + "--tmp_dir", + dest="tmp_dir", + type=_paths, + default=None, + help="Directory for intermediates (default: the system temporary directory).", + ) + + def _add_device(parser: argparse.ArgumentParser, download: bool = True) -> None: """Add the shared device / verbosity / download options to a sub-parser.""" device = parser.add_mutually_exclusive_group() @@ -109,6 +127,7 @@ def main() -> None: "e.g. --set iterations=300 (repeatable).", ) _add_device(reg) + _add_tmp_dir(reg) # eval ------------------------------------------------------------------- ev = subparsers.add_parser( @@ -145,6 +164,7 @@ def main() -> None: ) ev.add_argument("-o", "--output", type=_paths, default=Path("./Output").resolve(), help="Output directory.") _add_device(ev) + _add_tmp_dir(ev) # uncertainty ------------------------------------------------------------ unc = subparsers.add_parser( @@ -166,6 +186,7 @@ def main() -> None: ) unc.add_argument("-o", "--output", type=_paths, default=Path("./Output").resolve(), help="Output directory.") _add_device(unc) + _add_tmp_dir(unc) args = parser.parse_args() app = ImpactRegKonfAIApp( @@ -187,6 +208,7 @@ def main() -> None: tta=args.tta, keep_dvf=args.uncertainty, config_overrides=args.config_overrides, + tmp_dir=args.tmp_dir, ) elif args.command == "eval": @@ -214,6 +236,7 @@ def main() -> None: gpu=gpu, cpu=args.cpu, quiet=args.quiet, + tmp_dir=args.tmp_dir, ) elif args.command == "uncertainty": @@ -225,6 +248,7 @@ def main() -> None: gpu=gpu, cpu=args.cpu, quiet=args.quiet, + tmp_dir=args.tmp_dir, ) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 3e659a20..96249ff5 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -104,6 +104,28 @@ def _output_path(dest_dir: Path, stem: str, suffixes: str) -> Path: return dest_dir / (stem + suffixes) +def _work_dir(tmp_dir: Path | None, prefix: str) -> Path: + """A private scratch directory for one command's intermediates. + + Under ``tmp_dir`` when the caller named one, under ``tempfile.gettempdir()`` otherwise — which is + the same contract every other KonfAI app CLI offers through ``--tmp-dir``. + + THE DEFAULT IS NOT ALWAYS A GOOD PLACE, WHICH IS WHY THE OPTION EXISTS. What is staged here is + volume-sized: the moved image and the displacement field are written before being collected into + ``--output``. Where TMPDIR is a tmpfs that traffic is charged to RAM, on top of the volumes the run + already holds; on a large 3D case that is what fills memory or the temp quota mid-run. A caller who + knows better — a pipeline node with its results on real disk — names a directory here instead of + overriding ``TMPDIR`` from outside, which would also move things TMPDIR legitimately owns (torch's + DataLoader opens its worker sockets there, and AF_UNIX addresses cap at ~108 bytes). + + The directory returned is always freshly created and owned by the caller of this function, never + ``tmp_dir`` itself: the command removes what it made and leaves the directory it was given. + """ + if tmp_dir is not None: + tmp_dir.mkdir(parents=True, exist_ok=True) + return Path(tempfile.mkdtemp(prefix=prefix, dir=str(tmp_dir) if tmp_dir is not None else None)) + + def _copy_output(src: Path, dest_dir: Path, stem: str) -> Path: """Copy an output beside the results, keeping the form the preset produced (file or store).""" dest = _output_path(dest_dir, stem, "".join(src.suffixes)) @@ -194,6 +216,13 @@ def _infer_preset( command += ["-i", str(fixed_mask or _neutral_mask(work / "FixedMask.mha"))] command += ["-i", str(moving_mask or _neutral_mask(work / "MovingMask.mha"))] command += ["-o", str(out)] + # Hand konfai-apps a workspace we own, which is what every other app CLI does by exposing + # --tmp-dir. Without it konfai-apps auto-creates one under TMPDIR, writes the prediction to + # ./Predictions inside it, and copies that into `-o` before deleting it: one extra full-size + # write of the moved image AND the displacement field, on whatever filesystem TMPDIR names. + # Given a caller-owned workspace it writes straight into `-o` (see konfai_apps + # _stage_result_dir / _collect_result), so the copy disappears and `out` is where it always was. + command += ["--tmp-dir", str(out)] if tta: command += ["--tta", str(tta)] # Preset-parameter tuning: forwarded verbatim to `konfai-apps infer --set` (applies to every preset). @@ -228,11 +257,14 @@ def register( tta: int = 0, keep_dvf: bool = False, config_overrides: list[str] | None = None, + tmp_dir: Path | None = None, ) -> None: """Register each fixed/moving pair with the selected presets and ensemble their DVFs. Masks are optional and restrict the metric region; when omitted a whole-image mask is auto-filled, so every preset app always receives the four inputs (fixed, moving, fixed mask, moving mask) it declares. + + ``tmp_dir`` names where the intermediates are staged; see :func:`_work_dir`. """ for index, (fixed_image, moving_image) in enumerate(zip(fixed_images, moving_images, strict=True)): case_out = output / f"P{index:03d}" @@ -241,7 +273,7 @@ def register( # caller asks, so `uncertainty` can measure the ensemble spread afterwards. if keep_dvf: (case_out / _ENSEMBLE_DIR).mkdir(parents=True, exist_ok=True) - work = Path(tempfile.mkdtemp(prefix="impact_reg_")) + work = _work_dir(tmp_dir, "impact_reg_") try: # Masks are optional (they restrict the metric region); pass only those the caller gave and # let konfai-apps fill the rest with an all-ones default — no input read on the no-mask path. @@ -324,6 +356,7 @@ def evaluate( gpu: list[int] = [], cpu: int | None = None, quiet: bool = False, + tmp_dir: Path | None = None, ) -> None: """Evaluate a registration on any subset of modalities (image MAE, seg Dice, landmark TRE). @@ -337,7 +370,7 @@ def evaluate( transform_path = transforms[index] if index < len(transforms) else None transform = sitk.ReadTransform(str(transform_path)) if transform_path else sitk.Transform() eval_out = output / f"P{index:03d}" / "Evaluation" - work = Path(tempfile.mkdtemp(prefix="impact_reg_eval_")) + work = _work_dir(tmp_dir, "impact_reg_eval_") try: # Image: moving resampled onto the fixed grid vs fixed (MAE). Mask is optional. if index < len(fixed_images) and index < len(moving_images): @@ -409,6 +442,7 @@ def uncertainty( gpu: list[int] = [], cpu: int | None = None, quiet: bool = False, + tmp_dir: Path | None = None, ) -> None: """Estimate registration uncertainty as the voxel-wise spread of an ensemble of displacement fields. @@ -419,7 +453,7 @@ def uncertainty( """ if len(dvfs) < 2: raise ValueError("Uncertainty needs at least two ensemble displacement fields.") - work = Path(tempfile.mkdtemp(prefix="impact_reg_unc_")) + work = _work_dir(tmp_dir, "impact_reg_unc_") try: reference = read_displacement_field(dvfs[0]) rank = reference.GetDimension() From f6b2684b388fbc017cb4413608abab5026a1aa58 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 00:21:38 +0200 Subject: [PATCH 07/39] fix(impact-reg): forward the workspace to every nested konfai-apps command --tmp-dir reached konfai-apps infer but not the two other nested invocations, so a caller who placed the staging deliberately still had 'uncertainty' stage its Uncertainties under the system TMPDIR, and the three in-process 'evaluate' calls auto-create a workspace of their own. Both are the traffic the option exists to move, so the option only half worked. uncertainty now passes the private directory _work_dir already made inside the caller's tmp_dir, and evaluate forwards the same one through the tmp_dir the konfai-apps API accepts. Tests pin the forwarding for each nested command. --- .../impact_reg_konfai/impact_reg.py | 9 +- .../tests/unit/test_tmp_dir_forwarding.py | 91 +++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 96249ff5..560093a5 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -388,6 +388,7 @@ def evaluate( gpu=gpu, cpu=cpu, quiet=quiet, + tmp_dir=work, ) # Segmentation: moving seg warped onto fixed vs fixed seg (Dice). @@ -408,6 +409,7 @@ def evaluate( gpu=gpu, cpu=cpu, quiet=quiet, + tmp_dir=work, ) # Landmarks (TRE): the transform is defined on the fixed grid and maps fixed->moving, so the @@ -428,6 +430,7 @@ def evaluate( gpu=gpu, cpu=cpu, quiet=quiet, + tmp_dir=work, ) finally: shutil.rmtree(work, ignore_errors=True) @@ -470,7 +473,11 @@ def uncertainty( stack.SetDirection(direction.flatten()) sitk.WriteImage(stack, str(work / "DVFs.mha")) - command = ["konfai-apps", "uncertainty", _app_id(preset), "-i", str(work / "DVFs.mha"), "-o", str(output)] + # Same workspace hand-off as _infer_preset: without it konfai-apps auto-creates one under + # TMPDIR and stages Uncertainties there before copying it into -o, which is the staging + # this option exists to place. `work` is ours and already sits wherever tmp_dir asked for. + command = ["konfai-apps", "uncertainty", _app_id(preset), "-i", str(work / "DVFs.mha"), + "-o", str(output), "--tmp-dir", str(work)] if gpu: command += ["--gpu", *(str(g) for g in gpu)] elif cpu is not None: diff --git a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py new file mode 100644 index 00000000..be6f8b46 --- /dev/null +++ b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py @@ -0,0 +1,91 @@ +# 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 + +"""The workspace the orchestrator was given must reach the nested konfai-apps commands. + +``--tmp-dir`` exists so a caller whose system temporary directory is the wrong medium can place the +volume-sized staging itself. That only holds if EVERY nested invocation is told about it: a command +left without one auto-creates its own workspace under TMPDIR and stages there, which is precisely +the traffic the option was added to move. These tests pin the forwarding for each nested command, so +a new one cannot be added without it. +""" + +from pathlib import Path + +import pytest + +sitk = pytest.importorskip("SimpleITK") + +from impact_reg_konfai.impact_reg import ImpactRegKonfAIApp # noqa: E402 + + +def _tmp_dir_value(command: list[str]) -> str: + """The value ``--tmp-dir`` carries in a captured command line (fails the test when absent).""" + assert "--tmp-dir" in command, f"nested command carries no --tmp-dir: {command}" + return command[command.index("--tmp-dir") + 1] + + +def test_infer_preset_forwards_the_workspace(tmp_path: Path, monkeypatch, write_preset_output) -> None: + """``konfai-apps infer`` is told to work in the directory the orchestrator staged for it. + + Pointed at ``-o``, so konfai-apps writes the prediction straight there instead of staging it in a + throwaway workspace and copying it in. + """ + captured: list[list[str]] = [] + + def fake_run(command, **kwargs): + captured.append(list(command)) + # Stand in for the preset run: konfai-apps would leave its outputs under -o. + write_preset_output(Path(command[command.index("-o") + 1])) + return None + + monkeypatch.setattr("impact_reg_konfai.impact_reg.subprocess.run", fake_run) + + work = tmp_path / "work" + work.mkdir() + app = ImpactRegKonfAIApp() + app._infer_preset("FireANTs_SyN", tmp_path / "f.mha", tmp_path / "m.mha", None, None, work, [], None, True) + + assert len(captured) == 1 + assert _tmp_dir_value(captured[0]) == str(work / "FireANTs_SyN") + + +def test_uncertainty_forwards_the_workspace(tmp_path: Path, monkeypatch, write_preset_output) -> None: + """``konfai-apps uncertainty`` stages inside the caller's tmp_dir, not under the system TMPDIR.""" + captured: list[list[str]] = [] + monkeypatch.setattr( + "impact_reg_konfai.impact_reg.subprocess.run", + lambda command, **kwargs: captured.append(list(command)), + ) + + _, first = write_preset_output(tmp_path / "a") + _, second = write_preset_output(tmp_path / "b") + staging = tmp_path / "staging" + + ImpactRegKonfAIApp().uncertainty( + preset="FireANTs_SyN", + dvfs=[first, second], + output=tmp_path / "out", + quiet=True, + tmp_dir=staging, + ) + + assert len(captured) == 1 + # The workspace is the private directory _work_dir made INSIDE the caller's tmp_dir -- never the + # caller's own directory, which the command must leave standing. + forwarded = Path(_tmp_dir_value(captured[0])) + assert forwarded.parent == staging + assert forwarded != staging From a511cbc1f9169d611ed62f4b60dba60ebca05a92 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 00:40:37 +0200 Subject: [PATCH 08/39] feat(impact-reg): let a preset declare only its field, and derive the moved image A registration preset's output IS the displacement field; the moved image is that field applied to the moving. Requiring both made every preset carry a second output this layer can produce itself -- and for a tiled preset, blend it across every patch seam only for a caller that reads the field to discard it. _find_output gains required=False, so _infer_preset returns Moved or None, and register derives it when the preset left it out. Deriving it needs the moving image back, which is where the ensemble path was already wrong: it re-read with sitk.ReadImage, which cannot open a directory, so an ensemble over OME-Zarr inputs failed there and nowhere else -- the same regression 1.6.0 fixed for the single-preset path by never re-reading at all. Both paths now go through _read_image / _write_image, a dispatch mirroring konfai's read_displacement_field and the existing _write_displacement_field, so the format in is the format out and the resample never decides it. The resample itself goes through SimpleITK on the field's own grid, for the reason konfai's ResampleTransform gives: the stored displacement is in world (x, y, z) units and adding it onto a (z, y, x) voxel grid transposes the axes. --- CHANGELOG.md | 4 + .../impact_reg_konfai/impact_reg.py | 108 +++++++++++++++--- .../tests/unit/test_displacement_field_io.py | 29 +++++ .../tests/unit/test_orchestration.py | 33 ++++++ 4 files changed, 160 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e06c56a..195e60f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -152,6 +152,10 @@ written replaces it. a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the displacement field once per run instead of twice +- **impact-reg**: a preset may declare only its displacement field — `register` derives the moved image + from it rather than requiring every preset to write a second output. Reading the moving image now + handles an OME-Zarr store as well as an ITK file, which also fixes the ensemble path: averaging + several presets over OME-Zarr inputs failed there, and nowhere else, on `sitk.ReadImage` - **studio**: bundle icons through the app interface, and a way to stop Studio (#75) - **examples**: a Transform example -- a template folded out of a cohort, and drawn copies of a case diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 560093a5..b68fe6bf 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -17,8 +17,10 @@ """Orchestrator for IMPACT-Reg. Each IMPACT-Reg *preset* is a self-contained KonfAI app on ``VBoussot/ImpactReg`` (one preset = one app): -its model produces, on the FIXED grid, the moving image resampled onto the fixed image (``MovedImage``) -and the displacement field (``DisplacementField``). This orchestrator adds the registration-specific +its model produces the displacement field (``DisplacementField``) on the FIXED grid, and optionally the +moving image resampled onto the fixed one (``MovedImage``). The field is what a registration preset must +declare; the moved image IS that field applied to the moving, so a preset that leaves it out is complete +and this orchestrator derives it. This layer adds the registration-specific logic that does not fit the generic ``konfai-apps`` pipeline, split into three composable operations (mirroring ``konfai-apps`` infer/eval/uncertainty) so a UI/CLI can run them independently: @@ -79,15 +81,22 @@ def get_available_presets(force_update: bool = False) -> list[str]: return list(get_available_apps_on_hf_repo(IMPACT_REG_KONFAI_REPO, force_update)) -def _find_output(root: Path, stem: str) -> Path: +def _find_output(root: Path, stem: str, required: bool = True) -> Path | None: """Locate the single output named ``stem`` under ``root``, whatever form the preset wrote it in. + ``required=False`` returns None instead of raising, for an output a preset may legitimately not + declare: the displacement field is what a registration preset must produce, and the moved image is + that field applied to the moving -- derivable here (see :func:`_derive_moved`) rather than a second + thing every preset has to remember to write. + Matched on the name rather than on a fixed filename: a displacement field may come out as an ITK image or as an OME-Zarr store, and a store is a DIRECTORY whose ``Path.stem`` is "DVF.ome" -- so a fixed "DVF.mha" finds nothing and the run dies with the output sitting in the directory it listed. """ matches = sorted(root.rglob(f"{stem}.*")) if not matches: + if not required: + return None raise FileNotFoundError(f"Preset inference did not produce '{stem}' under {root}.") return matches[0] @@ -160,6 +169,74 @@ def _write_displacement_field(field: sitk.Image, dest: Path) -> None: sitk.WriteImage(field, str(dest)) +def _read_image(path: Path) -> sitk.Image: + """An image, from an ITK file OR an OME-Zarr store — the input side of :func:`_write_image`. + + ``sitk.ReadImage`` cannot open a store, so any path that re-reads an input has to know both forms + or silently only work for one. This is the one place that knows, mirroring konfai's + ``read_displacement_field`` for fields. + """ + if not path.is_dir(): + return sitk.ReadImage(str(path)) + from konfai.utils.dataset import data_to_image, ome_zarr_attributes + from konfai.utils.ome_zarr import get_ome_zarr_info, read_ome_zarr_data_slice + + axes = get_ome_zarr_info(path)["axes"] + n_axes = 1 + sum(axis in axes for axis in ("z", "y", "x")) # channel-first C[Z]YX + data, metadata = read_ome_zarr_data_slice(path, tuple(slice(None) for _ in range(n_axes))) + # Origin, Spacing and Direction together: NGFF scale/translation alone cannot express the + # direction matrix, so the geometry comes from the konfai sidecar through data_to_image. + return data_to_image(data, ome_zarr_attributes(metadata)) + + +def _write_image(image: sitk.Image, dest: Path) -> None: + """Write an image in the form ``dest`` names — the scalar counterpart of + :func:`_write_displacement_field`.""" + if "".join(dest.suffixes).endswith(".ome.zarr"): + from konfai.utils.dataset import image_to_data + from konfai.utils.ome_zarr import write_ome_zarr + + data, attributes = image_to_data(image) + write_ome_zarr( + dest, data, spacing=image.GetSpacing(), origin=image.GetOrigin(), attributes=dict(attributes) + ) + else: + sitk.WriteImage(image, str(dest)) + + +def _derive_moved(moving_image: Path, dvf_path: Path, dest_dir: Path, field: sitk.Image | None = None) -> Path: + """The moved image, resampled from the moving through the displacement field. + + A preset that emits only a field is complete: the moved image IS that field applied to the moving, + so deriving it belongs to the orchestrator rather than being a second output every preset has to + remember to declare. + + FORMAT IN, SAME FORMAT OUT. Both ends go through the dispatch above, so an OME-Zarr moving yields + an OME-Zarr moved and an ITK one an ITK file; the resample never decides the format. Reading the + moving with ``sitk.ReadImage`` instead -- what this replaces -- cannot open a store at all, which + is why the ensemble path could not be used with OME-Zarr inputs. + + Resampled through SimpleITK for the reason konfai's own ``ResampleTransform`` gives: the stored + displacement is in world (x, y, z) units, and adding it onto a (z, y, x) voxel grid by hand + transposes the axes and reads millimetres as voxels. The output grid is the field's own -- a + displacement field is defined ON the fixed grid, so that is where the moved image belongs. + """ + if field is None: + field = read_displacement_field(dvf_path) + # Read the grid off the field BEFORE the transform takes it: DisplacementFieldTransform assumes + # ownership of the image it is given and leaves it empty behind. + size, origin = field.GetSize(), field.GetOrigin() + spacing, direction = field.GetSpacing(), field.GetDirection() + transform = sitk.DisplacementFieldTransform(field) + moving = _read_image(moving_image) + moved = sitk.Resample( + moving, size, transform, sitk.sitkLinear, origin, spacing, direction, 0.0, moving.GetPixelID() + ) + dest = _output_path(dest_dir, "Moved", "".join(dvf_path.suffixes)) + _write_image(moved, dest) + return dest + + def _displacement_transform(dvf_path: Path) -> sitk.Transform: """Read a displacement field (3-component, fixed grid) as a SimpleITK transform.""" return sitk.DisplacementFieldTransform(read_displacement_field(dvf_path)) @@ -241,7 +318,7 @@ def _infer_preset( subprocess.run(command, check=True) # nosec B603 # The model emits both the moved image and the displacement field on the fixed grid; reusing them # (rather than re-resampling here) keeps the single-preset path free of any extra image read/write. - return _find_output(out, "Moved"), _find_output(out, "DVF") + return _find_output(out, "Moved", required=False), _find_output(out, "DVF") def register( self, @@ -301,23 +378,26 @@ def register( dvf_paths.append(dvf) if len(presets) == 1: - # One preset: the model already produced the moved image AND the displacement field on - # the fixed grid — reuse them verbatim. No input re-read, no re-resample, and the input - # format is whatever the model handled (OME-Zarr included). - _copy_output(moved_paths[0], case_out, "Moved") dvf_out = _copy_output(dvf_paths[0], case_out, "DVF") + if moved_paths[0] is not None: + # The model already produced the moved image on the fixed grid — reuse it + # verbatim. No input re-read, no re-resample, and the input format is whatever + # the model handled (OME-Zarr included). + _copy_output(moved_paths[0], case_out, "Moved") + else: + # A preset that declares only a field is complete: the moved image is that + # field applied to the moving, and producing it belongs here. + _derive_moved(moving_image, dvf_out, case_out) else: # Ensemble: average the presets' displacement fields (all on the fixed grid) and warp the # moving image once with that averaged field — the one output no single preset produced. avg_dvf = self._average_displacement(dvf_paths) dvf_out = _output_path(case_out, "DVF", "".join(dvf_paths[0].suffixes)) _write_displacement_field(avg_dvf, dvf_out) - transform = sitk.DisplacementFieldTransform(sitk.Cast(avg_dvf, sitk.sitkVectorFloat64)) - moving = sitk.ReadImage(str(moving_image)) - sitk.WriteImage( - sitk.Resample(moving, avg_dvf, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()), - str(_output_path(case_out, "Moved", ".mha")), - ) + # Through the same derivation as the single-preset path, which reads the moving in + # either form: the sitk.ReadImage this replaces cannot open a store at all, so an + # ensemble of OME-Zarr inputs failed here and nowhere else. + _derive_moved(moving_image, dvf_out, case_out, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64)) # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid displacement # field as a SimpleITK transform. diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index 288f5824..9ca93b94 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -35,6 +35,7 @@ _displacement_transform, _find_output, _output_path, + _read_image, _write_displacement_field, ) from konfai.utils.errors import TransformError # noqa: E402 @@ -215,3 +216,31 @@ def test_transform_reads_back_identically_from_either_form(tmp_path: Path, suffi reference = sitk.DisplacementFieldTransform(sitk.Image(original)) for point in ((9.0, -1.0, 12.0), (7.5, -2.5, 11.0)): assert restored.TransformPoint(point) == pytest.approx(reference.TransformPoint(point)) + + +@pytest.mark.skipif(not _zarr_v3_available(), reason="writing an OME-Zarr store needs zarr 3") +def test_read_image_opens_an_ome_zarr_store(tmp_path) -> None: + """An ordinary image reads back from a store, not only from an ITK file. + + ``sitk.ReadImage`` cannot open a directory, so every place that re-read an input was silently + ITK-only. That is what made the ensemble path unusable with OME-Zarr inputs while the + single-preset path — which reuses the model's output and never re-reads — worked fine. + """ + volume = np.arange(4 * 5 * 6, dtype=np.float32).reshape(4, 5, 6) + store = tmp_path / "moving.ome.zarr" + write_ome_zarr(store, volume[np.newaxis], spacing=(1.5, 2.0, 2.5), origin=(3.0, -1.0, 0.5)) + + image = _read_image(store) + + assert image.GetSpacing() == pytest.approx((1.5, 2.0, 2.5)) + assert image.GetOrigin() == pytest.approx((3.0, -1.0, 0.5)) + np.testing.assert_allclose(sitk.GetArrayFromImage(image), volume) + + +def test_read_image_still_opens_an_itk_file(tmp_path) -> None: + """The other half of the dispatch: a plain file goes straight through SimpleITK.""" + volume = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + path = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(volume), str(path)) + + np.testing.assert_allclose(sitk.GetArrayFromImage(_read_image(path)), volume) diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index 7bb2ee6f..898fc16e 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -148,3 +148,36 @@ def test_register_multi_preset_averages_and_warps_once(tmp_path: Path) -> None: # keep_dvf persists each preset's field for a later uncertainty pass assert (case / "Ensemble" / "A.mha").is_file() and (case / "Ensemble" / "B.mha").is_file() assert (case / "Moved.mha").is_file() + + +def test_register_derives_moved_when_the_preset_emits_only_a_field(tmp_path: Path) -> None: + """A preset is complete with a displacement field alone: the moved image is derived here. + + The field IS the registration; the moved image is that field applied to the moving. Requiring both + made every preset carry a second output whose content the orchestrator can produce itself — and + for the tiled presets, blend across every patch seam only to have it thrown away. + """ + moving = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.arange(8**3, dtype=np.float32).reshape(8, 8, 8)), str(moving)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + reference = sitk.ReadImage(str(moving)) + app = reg.ImpactRegKonfAIApp() + + def field_only(preset, fixed_image, mov_image, fixed_mask, moving_mask, work, *args, **kwargs): + out = Path(work) / preset + out.mkdir(parents=True, exist_ok=True) + _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) + return None, out / "DVF.mha" + + app._infer_preset = field_only # type: ignore[method-assign] + out = tmp_path / "Output" + app.register(["FireANTs_SyN"], [fixed], [moving], output=out) + + case = out / "P000" + assert (case / "Moved.mha").is_file(), "the orchestrator did not derive the moved image" + # moved(p) = moving(p + d), and d is +2 along x on a unit grid: moving is z*64 + y*8 + x. + moved = sitk.GetArrayFromImage(sitk.ReadImage(str(case / "Moved.mha"))) + np.testing.assert_allclose(moved[0, 0, 0], 2.0, atol=1e-6) + np.testing.assert_allclose(moved[1, 1, 0], 74.0, atol=1e-6) From f647c5655f38d9af746b0bf9b30a947ccd04d309 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 00:48:52 +0200 Subject: [PATCH 09/39] refactor(impact-reg): read and write through Dataset, not a two-format branch The dispatch this replaces knew exactly two formats -- an ITK file or an OME-Zarr store -- because those are the two whoever wrote it had in front of them. konfai already has the layer that knows them all: Dataset covers h5, DICOM, OME-Zarr and every ITK extension through SitkFile, which probes them itself, and re-detects a directory store from disk whatever the format token said. _dataset_entry addresses a bare path as a Dataset entry -- root is the parent, the stem is the entry, the case is empty -- so _read_image and _write_image are now two lines each and inherit every format konfai supports, in and out. evaluate's four input reads went through sitk.ReadImage as well, so evaluating an OME-Zarr pair failed the same way the ensemble path did. They go through the same reader now; no direct sitk.ReadImage of a caller-supplied input is left. --- .../impact_reg_konfai/impact_reg.py | 63 ++++++++++--------- 1 file changed, 33 insertions(+), 30 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index b68fe6bf..1d4101d5 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -38,6 +38,7 @@ import subprocess import tempfile from pathlib import Path +from typing import Any import numpy as np import SimpleITK as sitk @@ -169,39 +170,41 @@ def _write_displacement_field(field: sitk.Image, dest: Path) -> None: sitk.WriteImage(field, str(dest)) -def _read_image(path: Path) -> sitk.Image: - """An image, from an ITK file OR an OME-Zarr store — the input side of :func:`_write_image`. +def _dataset_entry(path: Path) -> tuple[Any, str, str]: + """Address one file or store as a konfai ``Dataset`` entry: ``(dataset, group, name)``. + + Dataset is the layer that already knows every format konfai supports -- h5, DICOM, OME-Zarr, and + every ITK extension through SitkFile, which probes them itself. Going through it is what makes + this orchestrator format-agnostic without owning a dispatch of its own: a hand-written + ``if path.is_dir()`` knows exactly the two formats whoever wrote it thought of. - ``sitk.ReadImage`` cannot open a store, so any path that re-reads an input has to know both forms - or silently only work for one. This is the one place that knows, mirroring konfai's - ``read_displacement_field`` for fields. + A dataset addresses ``{root}/{name}/{group}.{ext}``. An orchestrator input is a bare path, so the + case is empty: the parent directory is the root and the stem is the entry. The extension only + seeds the format token -- Dataset normalises it (``.ome.zarr`` / ``.zarr`` -> ``omezarr``) and + re-detects a directory store from disk regardless of what the token said. """ - if not path.is_dir(): - return sitk.ReadImage(str(path)) - from konfai.utils.dataset import data_to_image, ome_zarr_attributes - from konfai.utils.ome_zarr import get_ome_zarr_info, read_ome_zarr_data_slice + from konfai.utils.dataset import Dataset + + suffixes = "".join(path.suffixes) + stem = path.name[: len(path.name) - len(suffixes)] if suffixes else path.name + file_format = "omezarr" if suffixes.lower().endswith((".ome.zarr", ".zarr")) else suffixes.lstrip(".") + return Dataset(path.parent, file_format or "mha"), stem, "" - axes = get_ome_zarr_info(path)["axes"] - n_axes = 1 + sum(axis in axes for axis in ("z", "y", "x")) # channel-first C[Z]YX - data, metadata = read_ome_zarr_data_slice(path, tuple(slice(None) for _ in range(n_axes))) - # Origin, Spacing and Direction together: NGFF scale/translation alone cannot express the - # direction matrix, so the geometry comes from the konfai sidecar through data_to_image. - return data_to_image(data, ome_zarr_attributes(metadata)) + +def _read_image(path: Path) -> sitk.Image: + """An image, in whatever format it is stored — read through konfai's Dataset.""" + dataset, group, name = _dataset_entry(path) + return dataset.read_image(group, name) def _write_image(image: sitk.Image, dest: Path) -> None: - """Write an image in the form ``dest`` names — the scalar counterpart of - :func:`_write_displacement_field`.""" - if "".join(dest.suffixes).endswith(".ome.zarr"): - from konfai.utils.dataset import image_to_data - from konfai.utils.ome_zarr import write_ome_zarr + """Write an image in the form ``dest`` names — through the same layer that read it. - data, attributes = image_to_data(image) - write_ome_zarr( - dest, data, spacing=image.GetSpacing(), origin=image.GetOrigin(), attributes=dict(attributes) - ) - else: - sitk.WriteImage(image, str(dest)) + So the format out is the format in, for every format konfai writes, and not only for the two an + ``if`` here would have enumerated. + """ + dataset, group, name = _dataset_entry(dest) + dataset.write(group, name, image) def _derive_moved(moving_image: Path, dvf_path: Path, dest_dir: Path, field: sitk.Image | None = None) -> Path: @@ -454,10 +457,10 @@ def evaluate( try: # Image: moving resampled onto the fixed grid vs fixed (MAE). Mask is optional. if index < len(fixed_images) and index < len(moving_images): - fixed = sitk.ReadImage(str(fixed_images[index])) + fixed = _read_image(fixed_images[index]) moved = work / "moved_image.nii.gz" sitk.WriteImage( - sitk.Resample(sitk.ReadImage(str(moving_images[index])), fixed, transform), str(moved) + sitk.Resample(_read_image(moving_images[index]), fixed, transform), str(moved) ) app.evaluate( inputs=[[fixed_images[index]]], @@ -473,11 +476,11 @@ def evaluate( # Segmentation: moving seg warped onto fixed vs fixed seg (Dice). if index < len(gt_fixed_seg) and index < len(gt_moving_seg): - fixed_seg = sitk.ReadImage(str(gt_fixed_seg[index])) + fixed_seg = _read_image(gt_fixed_seg[index]) moved_seg = work / "moved_seg.nii.gz" sitk.WriteImage( sitk.Resample( - sitk.ReadImage(str(gt_moving_seg[index])), fixed_seg, transform, sitk.sitkNearestNeighbor + _read_image(gt_moving_seg[index]), fixed_seg, transform, sitk.sitkNearestNeighbor ), str(moved_seg), ) From c090453ae961e8bad902a6b84f94916dd0349458 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 01:31:28 +0200 Subject: [PATCH 10/39] refactor(impact-reg): the preset owes a field, and nothing else _infer_preset looked for a Moved as well as a DVF and reused it when it was there. That made the moved image half a contract: presets had to write it, the orchestrator had to branch on it, and a tiled preset blended a full-size one across every patch seam for a caller that then read only the field. A registration app produces a displacement field on the fixed grid, in whatever format it declares. That is the whole contract. Everything computable from it -- the moved image above all -- is this layer's job, so the moved is now always derived and the branch is gone. Measured while checking the output layout: a preset fed .mha inputs writes its field as .ome.zarr when that is what it declares. So the input's form says nothing about the output's; the derived moved follows the FIELD, which is the only thing the preset committed to. The docstring said otherwise. --- CHANGELOG.md | 10 ++-- .../impact_reg_konfai/impact_reg.py | 54 ++++++++----------- .../tests/unit/test_orchestration.py | 10 ++-- 3 files changed, 31 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 195e60f9..a4523300 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -152,10 +152,12 @@ written replaces it. a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the displacement field once per run instead of twice -- **impact-reg**: a preset may declare only its displacement field — `register` derives the moved image - from it rather than requiring every preset to write a second output. Reading the moving image now - handles an OME-Zarr store as well as an ITK file, which also fixes the ensemble path: averaging - several presets over OME-Zarr inputs failed there, and nowhere else, on `sitk.ReadImage` +- **impact-reg**: a registration preset now owes exactly one output, its displacement field, in whatever + format it declares — `register` derives the moved image from it instead of expecting a second output. + A preset can drop `MovedImage` entirely, which for a tiled one also drops blending a full-size moved + across every patch seam for a caller that has the field. Reading the moving image handles an OME-Zarr + store as well as an ITK file, which fixes the ensemble path too: averaging several presets over + OME-Zarr inputs failed there, and nowhere else, on `sitk.ReadImage` - **studio**: bundle icons through the app interface, and a way to stop Studio (#75) - **examples**: a Transform example -- a template folded out of a cohort, and drawn copies of a case diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 1d4101d5..92e8b4b5 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -17,10 +17,10 @@ """Orchestrator for IMPACT-Reg. Each IMPACT-Reg *preset* is a self-contained KonfAI app on ``VBoussot/ImpactReg`` (one preset = one app): -its model produces the displacement field (``DisplacementField``) on the FIXED grid, and optionally the -moving image resampled onto the fixed one (``MovedImage``). The field is what a registration preset must -declare; the moved image IS that field applied to the moving, so a preset that leaves it out is complete -and this orchestrator derives it. This layer adds the registration-specific +its model produces one thing: the displacement field (``DisplacementField``) on the FIXED grid, in +whatever format it declares. That is the whole contract. The moved image IS that field applied to the +moving, so this orchestrator derives it rather than asking every preset to write it as well. This layer +adds the registration-specific logic that does not fit the generic ``konfai-apps`` pipeline, split into three composable operations (mirroring ``konfai-apps`` infer/eval/uncertainty) so a UI/CLI can run them independently: @@ -82,22 +82,15 @@ def get_available_presets(force_update: bool = False) -> list[str]: return list(get_available_apps_on_hf_repo(IMPACT_REG_KONFAI_REPO, force_update)) -def _find_output(root: Path, stem: str, required: bool = True) -> Path | None: +def _find_output(root: Path, stem: str) -> Path: """Locate the single output named ``stem`` under ``root``, whatever form the preset wrote it in. - ``required=False`` returns None instead of raising, for an output a preset may legitimately not - declare: the displacement field is what a registration preset must produce, and the moved image is - that field applied to the moving -- derivable here (see :func:`_derive_moved`) rather than a second - thing every preset has to remember to write. - Matched on the name rather than on a fixed filename: a displacement field may come out as an ITK image or as an OME-Zarr store, and a store is a DIRECTORY whose ``Path.stem`` is "DVF.ome" -- so a fixed "DVF.mha" finds nothing and the run dies with the output sitting in the directory it listed. """ matches = sorted(root.rglob(f"{stem}.*")) if not matches: - if not required: - return None raise FileNotFoundError(f"Preset inference did not produce '{stem}' under {root}.") return matches[0] @@ -214,10 +207,11 @@ def _derive_moved(moving_image: Path, dvf_path: Path, dest_dir: Path, field: sit so deriving it belongs to the orchestrator rather than being a second output every preset has to remember to declare. - FORMAT IN, SAME FORMAT OUT. Both ends go through the dispatch above, so an OME-Zarr moving yields - an OME-Zarr moved and an ITK one an ITK file; the resample never decides the format. Reading the - moving with ``sitk.ReadImage`` instead -- what this replaces -- cannot open a store at all, which - is why the ensemble path could not be used with OME-Zarr inputs. + THE FORMAT FOLLOWS THE FIELD, not the moving. Measured: a preset fed ``.mha`` inputs writes its + field as ``.ome.zarr`` when that is what it declares, so the input's form says nothing about the + output's -- the field is the only thing the preset committed to, so the derived moved matches it. + The moving is read through the dataset layer either way; ``sitk.ReadImage``, what this replaces, + cannot open a store at all, which is why the ensemble path failed on OME-Zarr inputs. Resampled through SimpleITK for the reason konfai's own ``ResampleTransform`` gives: the stored displacement is in world (x, y, z) units, and adding it onto a (z, y, x) voxel grid by hand @@ -280,8 +274,8 @@ def _infer_preset( quiet: bool, tta: int = 0, config_overrides: list[str] | None = None, - ) -> tuple[Path, Path]: - """Run one preset app on the fixed/moving pair (+ optional masks); return its (moved, displacement) paths. + ) -> Path: + """Run one preset app on the fixed/moving pair (+ optional masks); return its displacement field. Each preset runs through the ``konfai-apps`` CLI in its own subprocess: konfai keeps process-global state (its ``Config`` singleton, the ``KONFAI_*`` environment), so several preset @@ -319,9 +313,12 @@ def _infer_preset( if self._force_update: command.append("--force_update") subprocess.run(command, check=True) # nosec B603 - # The model emits both the moved image and the displacement field on the fixed grid; reusing them - # (rather than re-resampling here) keeps the single-preset path free of any extra image read/write. - return _find_output(out, "Moved", required=False), _find_output(out, "DVF") + # THE PRESET'S CONTRACT IS THE FIELD, AND ONLY THE FIELD. A registration app produces a + # displacement field on the fixed grid, in whatever format it declares; anything else that can + # be computed from it -- the moved image above all -- is this layer's job. Looking for a Moved + # here would make every preset carry an output it does not owe, and a tiled one blend it across + # every patch seam for a caller that has the field. + return _find_output(out, "DVF") def register( self, @@ -360,9 +357,9 @@ def register( fixed_mask = fixed_masks[index] if index < len(fixed_masks) else None moving_mask = moving_masks[index] if index < len(moving_masks) else None - moved_paths, dvf_paths = [], [] + dvf_paths = [] for preset in presets: - moved, dvf = self._infer_preset( + dvf = self._infer_preset( preset, fixed_image, moving_image, @@ -377,20 +374,11 @@ def register( ) if keep_dvf: dvf = _copy_output(dvf, case_out / _ENSEMBLE_DIR, preset) - moved_paths.append(moved) dvf_paths.append(dvf) if len(presets) == 1: dvf_out = _copy_output(dvf_paths[0], case_out, "DVF") - if moved_paths[0] is not None: - # The model already produced the moved image on the fixed grid — reuse it - # verbatim. No input re-read, no re-resample, and the input format is whatever - # the model handled (OME-Zarr included). - _copy_output(moved_paths[0], case_out, "Moved") - else: - # A preset that declares only a field is complete: the moved image is that - # field applied to the moving, and producing it belongs here. - _derive_moved(moving_image, dvf_out, case_out) + _derive_moved(moving_image, dvf_out, case_out) else: # Ensemble: average the presets' displacement fields (all on the fixed grid) and warp the # moving image once with that averaged field — the one output no single preset produced. diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index 898fc16e..e02b8f52 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -98,21 +98,19 @@ def test_average_displacement_is_the_voxelwise_mean_with_reference_geometry(tmp_ def _stub_infer(app: reg.ImpactRegKonfAIApp, moving_image: Path, dvf_by_preset: dict[str, tuple]): - """Replace ``_infer_preset`` so it writes a Moved.mha + a constant DVF.mha (per preset) on the moving grid.""" + """Replace ``_infer_preset`` so it writes a constant DVF.mha per preset on the moving grid.""" reference = sitk.ReadImage(str(moving_image)) def fake(preset, fixed_image, mov_image, fixed_mask, moving_mask, work, *args, **kwargs): out = Path(work) / preset out.mkdir(parents=True, exist_ok=True) - moved = sitk.Image(reference) - sitk.WriteImage(moved, str(out / "Moved.mha")) _write_dvf(out / "DVF.mha", dvf_by_preset[preset], reference) - return out / "Moved.mha", out / "DVF.mha" + return out / "DVF.mha" app._infer_preset = fake # type: ignore[method-assign] -def test_register_single_preset_reuses_model_outputs(tmp_path: Path) -> None: +def test_register_single_preset_reuses_the_field_and_derives_the_moved(tmp_path: Path) -> None: moving = tmp_path / "moving.mha" sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) fixed = tmp_path / "fixed.mha" @@ -169,7 +167,7 @@ def field_only(preset, fixed_image, mov_image, fixed_mask, moving_mask, work, *a out = Path(work) / preset out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) - return None, out / "DVF.mha" + return out / "DVF.mha" app._infer_preset = field_only # type: ignore[method-assign] out = tmp_path / "Output" From 6bf61bd95c4b991d76250d2811003e51376495ad Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 01:41:22 +0200 Subject: [PATCH 11/39] fix(impact-reg): count the cases konfai-apps produced, not the arguments given register numbered its cases from the command line -- one P{index} per (fixed, moving) pair of paths -- while konfai-apps numbers them from the EXPANDED units: each -i is a group, a file or a store is one unit, and a plain directory is walked so every volume inside becomes its own. The two agree only while each group holds exactly one file. Hand it a directory and they diverge: konfai-apps registered every volume in it, _find_output took sorted(rglob(...))[0], and every case after the first was written to disk and dropped without a word. Verified on a real two-case run before the fix -- P001 computed, then discarded -- and after it: two directories in, two complete cases out, each with its field, its moved image and its transform. So the presets now run once over the whole cohort rather than once per pair (one model load instead of N), _find_outputs returns a field per case, and register iterates over the cases konfai-apps named rather than renumbering its own. The moving volume each case is derived from comes from _list_input_units, the same expansion konfai-apps uses, so the two orders cannot drift apart. --- CHANGELOG.md | 5 + .../impact_reg_konfai/impact_reg.py | 145 ++++++++++++------ .../tests/unit/test_displacement_field_io.py | 23 ++- .../tests/unit/test_orchestration.py | 16 +- .../tests/unit/test_tmp_dir_forwarding.py | 8 +- 5 files changed, 132 insertions(+), 65 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4523300..8ffaba14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -152,6 +152,11 @@ written replaces it. a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the displacement field once per run instead of twice +- **impact-reg**: `register` accepts a whole dataset per input, not only one volume — a directory is + expanded into one case per volume it holds, exactly as `konfai-apps infer` already does, and every case + gets its own field, moved image and transform. Previously the cases were counted from the command-line + arguments while konfai-apps counted them from the expanded units, so a directory input produced N + results and only the first was collected, silently - **impact-reg**: a registration preset now owes exactly one output, its displacement field, in whatever format it declares — `register` derives the moved image from it instead of expecting a second output. A preset can drop `MovedImage` entirely, which for a tiled one also drops blending a full-size moved diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 92e8b4b5..ad75bb3e 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -82,17 +82,23 @@ def get_available_presets(force_update: bool = False) -> list[str]: return list(get_available_apps_on_hf_repo(IMPACT_REG_KONFAI_REPO, force_update)) -def _find_output(root: Path, stem: str) -> Path: - """Locate the single output named ``stem`` under ``root``, whatever form the preset wrote it in. +def _find_outputs(root: Path, stem: str) -> dict[str, Path]: + """Every output named ``stem`` under ``root``, keyed by the CASE it belongs to. Matched on the name rather than on a fixed filename: a displacement field may come out as an ITK image or as an OME-Zarr store, and a store is a DIRECTORY whose ``Path.stem`` is "DVF.ome" -- so a fixed "DVF.mha" finds nothing and the run dies with the output sitting in the directory it listed. + + A MAPPING, NOT THE FIRST MATCH. konfai-apps writes one dataset per output group -- + ``///.`` -- and one run produces as many cases as the inputs expanded + to, which is not the number of paths on the command line: a directory is walked so each volume in + it becomes its own case. Taking ``matches[0]`` therefore kept P000 and dropped every case after it, + silently, with the results sitting on disk. The case is the entry's parent directory. """ matches = sorted(root.rglob(f"{stem}.*")) if not matches: raise FileNotFoundError(f"Preset inference did not produce '{stem}' under {root}.") - return matches[0] + return {match.parent.name: match for match in matches} def _output_path(dest_dir: Path, stem: str, suffixes: str) -> Path: @@ -107,6 +113,15 @@ def _output_path(dest_dir: Path, stem: str, suffixes: str) -> Path: return dest_dir / (stem + suffixes) +def _neutral_masks(work: Path, side: str, count: int) -> list[Path]: + """``count`` all-ones sentinels, one per case, standing in for a mask group the caller left out. + + Input groups pair by POSITION, so a group that is present at all must carry as many units as the + others: one sentinel would pair with the first case and leave the rest unmasked on that side. + """ + return [_neutral_mask(work / f"{side}Mask_{index:03d}.mha") for index in range(count)] + + def _work_dir(tmp_dir: Path | None, prefix: str) -> Path: """A private scratch directory for one command's intermediates. @@ -264,31 +279,40 @@ def __init__(self, download: bool = False, force_update: bool = False) -> None: def _infer_preset( self, preset: str, - fixed_image: Path, - moving_image: Path, - fixed_mask: Path | None, - moving_mask: Path | None, + fixed_images: list[Path], + moving_images: list[Path], + fixed_masks: list[Path], + moving_masks: list[Path], + n_cases: int, work: Path, gpu: list[int], cpu: int | None, quiet: bool, tta: int = 0, config_overrides: list[str] | None = None, - ) -> Path: - """Run one preset app on the fixed/moving pair (+ optional masks); return its displacement field. + ) -> dict[str, Path]: + """Run one preset app on every case at once; return its displacement field per case. + + ONE RUN, NOT ONE PER CASE. Each ``-i`` is an input GROUP, and konfai-apps expands each group's + paths into units -- a file is one, a store or DICOM series is one, a plain directory is walked + so each volume in it becomes one -- then pairs the groups by position into ``Dataset/P{i:03d}``. + So the whole cohort goes in a single invocation and the model is loaded once, rather than once + per case. + + Each preset still runs in its own subprocess: konfai keeps process-global state (its ``Config`` + singleton, the ``KONFAI_*`` environment), so several preset inferences in one process would clash. - Each preset runs through the ``konfai-apps`` CLI in its own subprocess: konfai keeps - process-global state (its ``Config`` singleton, the ``KONFAI_*`` environment), so several preset - inferences in one process would clash. The ``-i`` inputs map positionally to the app's input groups. Masks are optional: konfai-apps fills any we omit with an all-ones default, so with no mask we pass only fixed+moving. Because the mapping is positional, a lone moving mask still needs the fixed-mask - slot present, so send the pair (defaulting the absent one to an all-ones sentinel) once either is given. + slot present -- filled with one all-ones sentinel per case so the groups keep pairing. """ out = work / preset - command = ["konfai-apps", "infer", _app_id(preset), "-i", str(fixed_image), "-i", str(moving_image)] - if fixed_mask is not None or moving_mask is not None: - command += ["-i", str(fixed_mask or _neutral_mask(work / "FixedMask.mha"))] - command += ["-i", str(moving_mask or _neutral_mask(work / "MovingMask.mha"))] + command = ["konfai-apps", "infer", _app_id(preset)] + command += ["-i", *(str(path) for path in fixed_images)] + command += ["-i", *(str(path) for path in moving_images)] + if fixed_masks or moving_masks: + command += ["-i", *(str(path) for path in fixed_masks or _neutral_masks(work, "Fixed", n_cases))] + command += ["-i", *(str(path) for path in moving_masks or _neutral_masks(work, "Moving", n_cases))] command += ["-o", str(out)] # Hand konfai-apps a workspace we own, which is what every other app CLI does by exposing # --tmp-dir. Without it konfai-apps auto-creates one under TMPDIR, writes the prediction to @@ -318,7 +342,7 @@ def _infer_preset( # be computed from it -- the moved image above all -- is this layer's job. Looking for a Moved # here would make every preset carry an output it does not owe, and a tiled one blend it across # every patch seam for a caller that has the field. - return _find_output(out, "DVF") + return _find_outputs(out, "DVF") def register( self, @@ -336,42 +360,69 @@ def register( config_overrides: list[str] | None = None, tmp_dir: Path | None = None, ) -> None: - """Register each fixed/moving pair with the selected presets and ensemble their DVFs. + """Register every case with the selected presets and ensemble their DVFs. + + A case is whatever konfai-apps makes of the inputs: one image per group gives one case, and a + directory per group gives one case per volume inside it, paired by position. So the caller may + pass single volumes or whole datasets, exactly as ``konfai-apps infer`` accepts them. Masks are optional and restrict the metric region; when omitted a whole-image mask is auto-filled, so every preset app always receives the four inputs (fixed, moving, fixed mask, moving mask) it declares. ``tmp_dir`` names where the intermediates are staged; see :func:`_work_dir`. """ - for index, (fixed_image, moving_image) in enumerate(zip(fixed_images, moving_images, strict=True)): - case_out = output / f"P{index:03d}" - case_out.mkdir(parents=True, exist_ok=True) - # The per-preset displacement fields are large; only persist them (under Ensemble/) when the - # caller asks, so `uncertainty` can measure the ensemble spread afterwards. - if keep_dvf: - (case_out / _ENSEMBLE_DIR).mkdir(parents=True, exist_ok=True) - work = _work_dir(tmp_dir, "impact_reg_") - try: - # Masks are optional (they restrict the metric region); pass only those the caller gave and - # let konfai-apps fill the rest with an all-ones default — no input read on the no-mask path. - fixed_mask = fixed_masks[index] if index < len(fixed_masks) else None - moving_mask = moving_masks[index] if index < len(moving_masks) else None + # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into + # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs + # the groups by position. Asking it for the moving group's units is what tells this layer which + # volume belongs to which case, in the same order it will use, so a dataset in and a single pair + # in go down one path. + moving_units = [source for source, _ in KonfAIApp._list_input_units(list(moving_images))] + + work = _work_dir(tmp_dir, "impact_reg_") + try: + fields_by_preset = { + preset: self._infer_preset( + preset, + list(fixed_images), + list(moving_images), + list(fixed_masks), + list(moving_masks), + len(moving_units), + work, + gpu, + cpu, + quiet, + tta, + config_overrides, + ) + for preset in presets + } + cases = sorted(fields_by_preset[presets[0]]) + for preset, fields in fields_by_preset.items(): + if sorted(fields) != cases: + raise RuntimeError( + f"preset '{preset}' produced cases {sorted(fields)} where '{presets[0]}' produced " + f"{cases}; an ensemble can only be averaged case by case." + ) + if len(cases) != len(moving_units): + raise RuntimeError( + f"the presets produced {len(cases)} case(s) for {len(moving_units)} moving unit(s); " + "the moved image is derived per case and needs the two to line up." + ) + + for case, moving_image in zip(cases, moving_units, strict=True): + # konfai-apps already named the cases; reusing its names keeps the two layers' notion of + # a case identical instead of renumbering from the command line and hoping they agree. + case_out = output / case + case_out.mkdir(parents=True, exist_ok=True) + # The per-preset displacement fields are large; only persist them (under Ensemble/) when the + # caller asks, so `uncertainty` can measure the ensemble spread afterwards. + if keep_dvf: + (case_out / _ENSEMBLE_DIR).mkdir(parents=True, exist_ok=True) dvf_paths = [] for preset in presets: - dvf = self._infer_preset( - preset, - fixed_image, - moving_image, - fixed_mask, - moving_mask, - work, - gpu, - cpu, - quiet, - tta, - config_overrides, - ) + dvf = fields_by_preset[preset][case] if keep_dvf: dvf = _copy_output(dvf, case_out / _ENSEMBLE_DIR, preset) dvf_paths.append(dvf) @@ -393,8 +444,8 @@ def register( # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid displacement # field as a SimpleITK transform. sitk.WriteTransform(_displacement_transform(dvf_out), str(case_out / "Transform.h5")) - finally: - shutil.rmtree(work, ignore_errors=True) + finally: + shutil.rmtree(work, ignore_errors=True) def _average_displacement(self, dvf_paths: list[Path]) -> sitk.Image: """Average several presets' displacement fields (all on the same fixed grid) into one field. diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index 9ca93b94..e58c962e 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -33,7 +33,7 @@ from impact_reg_konfai.impact_reg import ( # noqa: E402 _copy_output, _displacement_transform, - _find_output, + _find_outputs, _output_path, _read_image, _write_displacement_field, @@ -71,18 +71,27 @@ def _write_store(dest: Path, field: "sitk.Image") -> Path: @pytest.mark.parametrize("suffix", [".mha", ".ome.zarr"]) -def test_find_output_locates_either_form(tmp_path: Path, suffix: str) -> None: +def test_find_outputs_locates_either_form(tmp_path: Path, suffix: str) -> None: """Discovery is by name, not by filename: a store is a directory whose stem is 'DVF.ome'.""" produced = tmp_path / "P000" produced.mkdir() _write_displacement_field(_field(), produced / f"DVF{suffix}") - assert _find_output(tmp_path, "DVF").name == f"DVF{suffix}" + assert _find_outputs(tmp_path, "DVF")["P000"].name == f"DVF{suffix}" -def test_find_output_reports_the_name_it_looked_for(tmp_path: Path) -> None: +def test_find_outputs_keys_every_case_the_run_produced(tmp_path: Path) -> None: + """One run yields one entry per case, so none is silently dropped.""" + for case in ("P000", "P001", "P002"): + (tmp_path / case).mkdir() + _write_displacement_field(_field(), tmp_path / case / "DVF.mha") + + assert sorted(_find_outputs(tmp_path, "DVF")) == ["P000", "P001", "P002"] + + +def test_find_outputs_reports_the_name_it_looked_for(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError, match="DVF"): - _find_output(tmp_path, "DVF") + _find_outputs(tmp_path, "DVF") @pytest.mark.parametrize("suffix", [".mha", ".ome.zarr"]) @@ -180,7 +189,7 @@ def test_output_path_clears_the_stem_whatever_the_previous_form(tmp_path: Path) def test_rerunning_in_the_other_form_leaves_one_output(tmp_path: Path) -> None: """Discovery is by stem, so a run that emits the other form must not leave the previous one - beside it -- ``_find_output`` returns the first match, and ``DVF.mha`` sorts before its store.""" + beside it -- discovery is by stem, and ``DVF.mha`` sorts before its store.""" source, destination = tmp_path / "src", tmp_path / "out" source.mkdir() destination.mkdir() @@ -191,7 +200,7 @@ def test_rerunning_in_the_other_form_leaves_one_output(tmp_path: Path) -> None: _copy_output(source / "DVF.ome.zarr", destination, "DVF") assert [p.name for p in destination.iterdir()] == ["DVF.ome.zarr"] - assert _find_output(destination, "DVF").name == "DVF.ome.zarr" + assert _find_outputs(destination, "DVF")[destination.name].name == "DVF.ome.zarr" def test_store_written_by_the_orchestrator_is_a_declared_field(tmp_path: Path) -> None: diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index e02b8f52..ae33259b 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -53,9 +53,9 @@ def test_get_available_presets_keeps_only_registration_apps(tmp_path: Path, monk assert reg.get_available_presets() == ["FireANTs_SyN", "Generic_Rigid"] -def test_find_output_raises_when_missing(tmp_path: Path) -> None: +def test_find_outputs_raises_when_missing(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError, match=r"Moved\.mha"): - reg._find_output(tmp_path, "Moved.mha") + reg._find_outputs(tmp_path, "Moved.mha") # --------------------------------------------------------------------------- mask sentinel @@ -101,11 +101,11 @@ def _stub_infer(app: reg.ImpactRegKonfAIApp, moving_image: Path, dvf_by_preset: """Replace ``_infer_preset`` so it writes a constant DVF.mha per preset on the moving grid.""" reference = sitk.ReadImage(str(moving_image)) - def fake(preset, fixed_image, mov_image, fixed_mask, moving_mask, work, *args, **kwargs): - out = Path(work) / preset + def fake(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + out = Path(work) / preset / "P000" out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", dvf_by_preset[preset], reference) - return out / "DVF.mha" + return {"P000": out / "DVF.mha"} app._infer_preset = fake # type: ignore[method-assign] @@ -163,11 +163,11 @@ def test_register_derives_moved_when_the_preset_emits_only_a_field(tmp_path: Pat reference = sitk.ReadImage(str(moving)) app = reg.ImpactRegKonfAIApp() - def field_only(preset, fixed_image, mov_image, fixed_mask, moving_mask, work, *args, **kwargs): - out = Path(work) / preset + def field_only(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + out = Path(work) / preset / "P000" out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) - return out / "DVF.mha" + return {"P000": out / "DVF.mha"} app._infer_preset = field_only # type: ignore[method-assign] out = tmp_path / "Output" diff --git a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py index be6f8b46..128636d5 100644 --- a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py +++ b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py @@ -48,8 +48,8 @@ def test_infer_preset_forwards_the_workspace(tmp_path: Path, monkeypatch, write_ def fake_run(command, **kwargs): captured.append(list(command)) - # Stand in for the preset run: konfai-apps would leave its outputs under -o. - write_preset_output(Path(command[command.index("-o") + 1])) + # Stand in for the preset run: konfai-apps leaves one case directory per unit under -o. + write_preset_output(Path(command[command.index("-o") + 1]) / "P000") return None monkeypatch.setattr("impact_reg_konfai.impact_reg.subprocess.run", fake_run) @@ -57,7 +57,9 @@ def fake_run(command, **kwargs): work = tmp_path / "work" work.mkdir() app = ImpactRegKonfAIApp() - app._infer_preset("FireANTs_SyN", tmp_path / "f.mha", tmp_path / "m.mha", None, None, work, [], None, True) + app._infer_preset( + "FireANTs_SyN", [tmp_path / "f.mha"], [tmp_path / "m.mha"], [], [], 1, work, [], None, True + ) assert len(captured) == 1 assert _tmp_dir_value(captured[0]) == str(work / "FireANTs_SyN") From 668ca09b111cd62b98d03957c0f87fd58ec51b4b Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 02:05:55 +0200 Subject: [PATCH 12/39] fix(impact-reg): expand eval's inputs the way register's are evaluate counted its cases the way register used to -- max() over the argument lists -- so a directory of volumes evaluated its first entry and nothing else, the same silent truncation just fixed one function above. Every group now goes through _units, the shared wrapper over konfai-apps' _list_input_units, so the two commands agree on what a case is. Transforms and landmark files expand too: .h5, .fcsv and .itk.txt are all supported extensions, so a directory of transforms pairs with a directory of volumes. --- CHANGELOG.md | 2 +- .../impact_reg_konfai/impact_reg.py | 22 ++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ffaba14..8f5408dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -152,7 +152,7 @@ written replaces it. a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the displacement field once per run instead of twice -- **impact-reg**: `register` accepts a whole dataset per input, not only one volume — a directory is +- **impact-reg**: `register` and `eval` accept a whole dataset per input, not only one volume — a directory is expanded into one case per volume it holds, exactly as `konfai-apps infer` already does, and every case gets its own field, moved image and transform. Previously the cases were counted from the command-line arguments while konfai-apps counted them from the expanded units, so a directory input produced N diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index ad75bb3e..7cedc101 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -113,6 +113,18 @@ def _output_path(dest_dir: Path, stem: str, suffixes: str) -> Path: return dest_dir / (stem + suffixes) +def _units(paths: list[Path]) -> list[Path]: + """What an input group expands to, in konfai-apps' own order. + + A file, an OME-Zarr store or a DICOM series is one unit; a plain directory is walked so every + supported file inside becomes one, sorted so groups pair consistently. Asking konfai-apps rather + than counting the command line is what keeps this layer's notion of a case identical to the one the + ``Dataset/P{i:03d}`` staging is built with -- counting arguments instead is how a directory input + came to register N volumes and report one. + """ + return [source for source, _ in KonfAIApp._list_input_units(list(paths))] if paths else [] + + def _neutral_masks(work: Path, side: str, count: int) -> list[Path]: """``count`` all-ones sentinels, one per case, standing in for a mask group the caller left out. @@ -376,7 +388,7 @@ def register( # the groups by position. Asking it for the moving group's units is what tells this layer which # volume belongs to which case, in the same order it will use, so a dataset in and a single pair # in go down one path. - moving_units = [source for source, _ in KonfAIApp._list_input_units(list(moving_images))] + moving_units = _units(moving_images) work = _work_dir(tmp_dir, "impact_reg_") try: @@ -487,6 +499,14 @@ def evaluate( already registered and only resampled onto the fixed grid (identity). """ app = KonfAIApp(_app_id(preset), self._download, self._force_update) + # Every group is expanded the same way ``register`` expands its inputs, so a directory of + # volumes evaluates case by case instead of collapsing to its first entry. Transforms and + # landmark files expand too: .h5, .fcsv and .itk.txt are all supported extensions. + fixed_images, moving_images = _units(fixed_images), _units(moving_images) + gt_fixed_seg, gt_moving_seg = _units(gt_fixed_seg), _units(gt_moving_seg) + gt_fixed_fid, gt_moving_fid = _units(gt_fixed_fid), _units(gt_moving_fid) + transforms = _units(transforms) + mask = _units(mask) if mask else mask n_cases = max(len(fixed_images), len(gt_fixed_seg), len(gt_fixed_fid)) for index in range(n_cases): transform_path = transforms[index] if index < len(transforms) else None From e8b18e59369b9fb8c21f10d86505c2e89c7a42c6 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 03:19:24 +0200 Subject: [PATCH 13/39] refactor(impact-reg): read through Dataset by building the dataset it needs Reading went through Dataset already, but by pretending a bare path was one: the parent directory taken for a root, the stem for a group, the case left empty. That is not a dataset, and it showed -- detection probes a case's entries, so on a flat layout it descended INTO the store and found nothing, forcing a format token to be guessed from the suffix, which is the dispatch Dataset was meant to replace. For a single-store backend like h5 it was not merely awkward but wrong: the file is the dataset, not an entry in one. A dataset is a root of cases holding groups, and it is built. So the path is linked into that layout first -- one case, one group, exactly as konfai-apps stages its own inputs -- and read with NO format named. konfai has a case to probe and detects the backend itself. The write side needed no such thing: // already IS that layout, so it goes through Dataset.write with the format named, which writing does require since nothing is on disk yet to detect. Verified: the OME-Zarr and ITK reader tests pass with no token, and a real two-case run over two directories still yields both cases complete. --- .../impact_reg_konfai/impact_reg.py | 80 ++++++++++--------- .../tests/unit/test_displacement_field_io.py | 17 ++-- 2 files changed, 54 insertions(+), 43 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 7cedc101..ab60e8b8 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -38,7 +38,6 @@ import subprocess import tempfile from pathlib import Path -from typing import Any import numpy as np import SimpleITK as sitk @@ -53,6 +52,10 @@ _ENSEMBLE_DIR = "Ensemble" +# Writing needs a format named -- nothing is on disk yet to detect one from. Only the store spellings +# need translating; every other suffix is already the token Dataset normalises (".mha" -> "mha"). +_FORMATS = {".ome.zarr": "omezarr", ".zarr": "omezarr"} + def _app_id(preset: str) -> str: """Resolve a preset to a KonfAIApp id: a local ``/`` path, or ``:`` on HF.""" @@ -190,44 +193,44 @@ def _write_displacement_field(field: sitk.Image, dest: Path) -> None: sitk.WriteImage(field, str(dest)) -def _dataset_entry(path: Path) -> tuple[Any, str, str]: - """Address one file or store as a konfai ``Dataset`` entry: ``(dataset, group, name)``. +def _read_image(path: Path, work: Path) -> sitk.Image: + """A volume, read through ``Dataset`` — after building the dataset it needs. - Dataset is the layer that already knows every format konfai supports -- h5, DICOM, OME-Zarr, and - every ITK extension through SitkFile, which probes them itself. Going through it is what makes - this orchestrator format-agnostic without owning a dispatch of its own: a hand-written - ``if path.is_dir()`` knows exactly the two formats whoever wrote it thought of. + A dataset is a root of CASES holding GROUPS. A path on the command line is neither, so it is linked + into that layout first, exactly as konfai-apps stages its own inputs. Pretending instead that the + parent directory is a root and the file a group -- what this used to do -- makes detection probe one + level too deep inside a store, and is simply wrong for a single-store backend like h5, where the + file IS the dataset rather than an entry in one. - A dataset addresses ``{root}/{name}/{group}.{ext}``. An orchestrator input is a bare path, so the - case is empty: the parent directory is the root and the stem is the entry. The extension only - seeds the format token -- Dataset normalises it (``.ome.zarr`` / ``.zarr`` -> ``omezarr``) and - re-detects a directory store from disk regardless of what the token said. + Built properly, no format is named: konfai has a case to probe and detects the backend itself, which + is why this reads an ITK file, an OME-Zarr store or anything else konfai supports without a branch + here enumerating them. """ from konfai.utils.dataset import Dataset - suffixes = "".join(path.suffixes) - stem = path.name[: len(path.name) - len(suffixes)] if suffixes else path.name - file_format = "omezarr" if suffixes.lower().endswith((".ome.zarr", ".zarr")) else suffixes.lstrip(".") - return Dataset(path.parent, file_format or "mha"), stem, "" - - -def _read_image(path: Path) -> sitk.Image: - """An image, in whatever format it is stored — read through konfai's Dataset.""" - dataset, group, name = _dataset_entry(path) - return dataset.read_image(group, name) + root = Path(tempfile.mkdtemp(prefix="entry_", dir=str(work))) + case = root / "P000" + case.mkdir() + (case / f"Entry{''.join(path.suffixes)}").symlink_to(path.resolve()) + return Dataset(root, "").read_image("Entry", "P000") -def _write_image(image: sitk.Image, dest: Path) -> None: - """Write an image in the form ``dest`` names — through the same layer that read it. +def _write_case_entry(image: sitk.Image, case_out: Path, group: str, file_format: str) -> None: + """Write one group of one case, through ``Dataset``. - So the format out is the format in, for every format konfai writes, and not only for the two an - ``if`` here would have enumerated. + THIS side really is a dataset, and always was: ``//.`` is a root of cases + holding groups, which is why konfai can read it back without being told a format. So the write goes + through the layer that owns that layout instead of composing the path by hand -- and writing, unlike + reading, does need a format named, because there is nothing on disk yet to detect one from. """ - dataset, group, name = _dataset_entry(dest) - dataset.write(group, name, image) + from konfai.utils.dataset import Dataset + + Dataset(case_out.parent, file_format).write(group, case_out.name, image) -def _derive_moved(moving_image: Path, dvf_path: Path, dest_dir: Path, field: sitk.Image | None = None) -> Path: +def _derive_moved( + moving_image: Path, dvf_path: Path, dest_dir: Path, work: Path, field: sitk.Image | None = None +) -> Path: """The moved image, resampled from the moving through the displacement field. A preset that emits only a field is complete: the moved image IS that field applied to the moving, @@ -252,12 +255,13 @@ def _derive_moved(moving_image: Path, dvf_path: Path, dest_dir: Path, field: sit size, origin = field.GetSize(), field.GetOrigin() spacing, direction = field.GetSpacing(), field.GetDirection() transform = sitk.DisplacementFieldTransform(field) - moving = _read_image(moving_image) + moving = _read_image(moving_image, work) moved = sitk.Resample( moving, size, transform, sitk.sitkLinear, origin, spacing, direction, 0.0, moving.GetPixelID() ) - dest = _output_path(dest_dir, "Moved", "".join(dvf_path.suffixes)) - _write_image(moved, dest) + suffixes = "".join(dvf_path.suffixes) + dest = _output_path(dest_dir, "Moved", suffixes) + _write_case_entry(moved, dest_dir, "Moved", _FORMATS.get(suffixes.lower(), suffixes.lstrip("."))) return dest @@ -441,7 +445,7 @@ def register( if len(presets) == 1: dvf_out = _copy_output(dvf_paths[0], case_out, "DVF") - _derive_moved(moving_image, dvf_out, case_out) + _derive_moved(moving_image, dvf_out, case_out, work) else: # Ensemble: average the presets' displacement fields (all on the fixed grid) and warp the # moving image once with that averaged field — the one output no single preset produced. @@ -451,7 +455,9 @@ def register( # Through the same derivation as the single-preset path, which reads the moving in # either form: the sitk.ReadImage this replaces cannot open a store at all, so an # ensemble of OME-Zarr inputs failed here and nowhere else. - _derive_moved(moving_image, dvf_out, case_out, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64)) + _derive_moved( + moving_image, dvf_out, case_out, work, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64) + ) # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid displacement # field as a SimpleITK transform. @@ -516,10 +522,10 @@ def evaluate( try: # Image: moving resampled onto the fixed grid vs fixed (MAE). Mask is optional. if index < len(fixed_images) and index < len(moving_images): - fixed = _read_image(fixed_images[index]) + fixed = _read_image(fixed_images[index], work) moved = work / "moved_image.nii.gz" sitk.WriteImage( - sitk.Resample(_read_image(moving_images[index]), fixed, transform), str(moved) + sitk.Resample(_read_image(moving_images[index], work), fixed, transform), str(moved) ) app.evaluate( inputs=[[fixed_images[index]]], @@ -535,11 +541,11 @@ def evaluate( # Segmentation: moving seg warped onto fixed vs fixed seg (Dice). if index < len(gt_fixed_seg) and index < len(gt_moving_seg): - fixed_seg = _read_image(gt_fixed_seg[index]) + fixed_seg = _read_image(gt_fixed_seg[index], work) moved_seg = work / "moved_seg.nii.gz" sitk.WriteImage( sitk.Resample( - _read_image(gt_moving_seg[index]), fixed_seg, transform, sitk.sitkNearestNeighbor + _read_image(gt_moving_seg[index], work), fixed_seg, transform, sitk.sitkNearestNeighbor ), str(moved_seg), ) diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index e58c962e..d2854b0c 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -231,15 +231,18 @@ def test_transform_reads_back_identically_from_either_form(tmp_path: Path, suffi def test_read_image_opens_an_ome_zarr_store(tmp_path) -> None: """An ordinary image reads back from a store, not only from an ITK file. - ``sitk.ReadImage`` cannot open a directory, so every place that re-read an input was silently - ITK-only. That is what made the ensemble path unusable with OME-Zarr inputs while the - single-preset path — which reuses the model's output and never re-reads — worked fine. + Read through ``Dataset`` with no format named: the path is staged into a real case/group layout + first, so konfai detects the backend itself. ``sitk.ReadImage`` cannot open a directory, which is + what made the ensemble path unusable with OME-Zarr inputs while the single-preset path — which + reuses the model's output and never re-reads — worked fine. """ volume = np.arange(4 * 5 * 6, dtype=np.float32).reshape(4, 5, 6) store = tmp_path / "moving.ome.zarr" write_ome_zarr(store, volume[np.newaxis], spacing=(1.5, 2.0, 2.5), origin=(3.0, -1.0, 0.5)) - image = _read_image(store) + work = tmp_path / "work" + work.mkdir() + image = _read_image(store, work) assert image.GetSpacing() == pytest.approx((1.5, 2.0, 2.5)) assert image.GetOrigin() == pytest.approx((3.0, -1.0, 0.5)) @@ -247,9 +250,11 @@ def test_read_image_opens_an_ome_zarr_store(tmp_path) -> None: def test_read_image_still_opens_an_itk_file(tmp_path) -> None: - """The other half of the dispatch: a plain file goes straight through SimpleITK.""" + """The same call, on a plain ITK file: one reader, no branch on the form.""" volume = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) path = tmp_path / "moving.mha" sitk.WriteImage(sitk.GetImageFromArray(volume), str(path)) - np.testing.assert_allclose(sitk.GetArrayFromImage(_read_image(path)), volume) + work = tmp_path / "work" + work.mkdir() + np.testing.assert_allclose(sitk.GetArrayFromImage(_read_image(path, work)), volume) From ef702f6533e492ef770dc74faf6cb50e3848924c Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Wed, 5 Aug 2026 03:36:10 +0200 Subject: [PATCH 14/39] feat(impact-reg): let a caller ask for the fields and nothing else The moved image and Transform.h5 are both derived FROM the displacement field -- a full-size resample of the moving and a full-size rewrite of the same voxels. Worth it for a caller that wants a registration to look at; waste for one that composes the field with another and derives its own moved from the total. The ExaSPIM tiled refinement is the second kind: it reads the field out of P000, composes it with the global pass and resamples through the COMPOSED transform, so the moved derived from the tiled residual alone is meaningless to it and the transform is a copy of a field it already holds. Both were written and deleted. --fields-only stops after the field. Nothing changes without it. --- CHANGELOG.md | 4 +++ apps/impact_reg/impact_reg_konfai/cli.py | 9 +++++ .../impact_reg_konfai/impact_reg.py | 33 ++++++++++++------- .../tests/unit/test_orchestration.py | 23 +++++++++++++ 4 files changed, 58 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f5408dc..d9cdd517 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -157,6 +157,10 @@ written replaces it. gets its own field, moved image and transform. Previously the cases were counted from the command-line arguments while konfai-apps counted them from the expanded units, so a directory input produced N results and only the first was collected, silently +- **impact-reg**: `register --fields-only` writes the displacement fields and stops there. The moved + image and `Transform.h5` are both derived from the field, at the cost of a full-size resample and a + full-size rewrite, so a caller that composes the field with its own and derives its own moved — the + ExaSPIM tiled refinement does exactly that — no longer pays for two outputs it deletes - **impact-reg**: a registration preset now owes exactly one output, its displacement field, in whatever format it declares — `register` derives the moved image from it instead of expecting a second output. A preset can drop `MovedImage` entirely, which for a tiled one also drops blending a full-size moved diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index 23e618a0..bdb425d9 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -126,6 +126,14 @@ def main() -> None: help="Tune a preset parameter, forwarded to 'konfai-apps infer --set' (applies to every preset), " "e.g. --set iterations=300 (repeatable).", ) + reg.add_argument( + "--fields-only", + "--fields_only", + dest="fields_only", + action="store_true", + help="Write the displacement fields only: skip the moved image and Transform.h5, both derived " + "from the field. For a caller that composes the field itself and would delete them.", + ) _add_device(reg) _add_tmp_dir(reg) @@ -209,6 +217,7 @@ def main() -> None: keep_dvf=args.uncertainty, config_overrides=args.config_overrides, tmp_dir=args.tmp_dir, + fields_only=args.fields_only, ) elif args.command == "eval": diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index ab60e8b8..9fbd0eb5 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -375,6 +375,7 @@ def register( keep_dvf: bool = False, config_overrides: list[str] | None = None, tmp_dir: Path | None = None, + fields_only: bool = False, ) -> None: """Register every case with the selected presets and ensemble their DVFs. @@ -386,6 +387,12 @@ def register( so every preset app always receives the four inputs (fixed, moving, fixed mask, moving mask) it declares. ``tmp_dir`` names where the intermediates are staged; see :func:`_work_dir`. + + ``fields_only`` writes the displacement fields and stops there. The moved image and + ``Transform.h5`` are both derived FROM the field, at the cost of a full-size resample and a + full-size rewrite -- worth it for a caller that wants a registration to look at, waste for one + that composes the field with another and derives its own. A caller that reads only the fields + should be able to say so rather than pay for outputs it deletes. """ # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs @@ -445,23 +452,27 @@ def register( if len(presets) == 1: dvf_out = _copy_output(dvf_paths[0], case_out, "DVF") - _derive_moved(moving_image, dvf_out, case_out, work) + if not fields_only: + _derive_moved(moving_image, dvf_out, case_out, work) else: # Ensemble: average the presets' displacement fields (all on the fixed grid) and warp the # moving image once with that averaged field — the one output no single preset produced. avg_dvf = self._average_displacement(dvf_paths) dvf_out = _output_path(case_out, "DVF", "".join(dvf_paths[0].suffixes)) _write_displacement_field(avg_dvf, dvf_out) - # Through the same derivation as the single-preset path, which reads the moving in - # either form: the sitk.ReadImage this replaces cannot open a store at all, so an - # ensemble of OME-Zarr inputs failed here and nowhere else. - _derive_moved( - moving_image, dvf_out, case_out, work, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64) - ) - - # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid displacement - # field as a SimpleITK transform. - sitk.WriteTransform(_displacement_transform(dvf_out), str(case_out / "Transform.h5")) + if not fields_only: + # Through the same derivation as the single-preset path, which reads the moving + # in either form: the sitk.ReadImage this replaces cannot open a store at all, + # so an ensemble of OME-Zarr inputs failed here and nowhere else. + _derive_moved( + moving_image, dvf_out, case_out, work, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64) + ) + + if not fields_only: + # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid + # displacement field as a SimpleITK transform. Another full-size write of the same + # voxels, which is why it goes with the moved image rather than being unconditional. + sitk.WriteTransform(_displacement_transform(dvf_out), str(case_out / "Transform.h5")) finally: shutil.rmtree(work, ignore_errors=True) diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index ae33259b..8ad4d7ce 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -179,3 +179,26 @@ def field_only(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, moved = sitk.GetArrayFromImage(sitk.ReadImage(str(case / "Moved.mha"))) np.testing.assert_allclose(moved[0, 0, 0], 2.0, atol=1e-6) np.testing.assert_allclose(moved[1, 1, 0], 74.0, atol=1e-6) + + +def test_register_fields_only_writes_nothing_derived(tmp_path: Path) -> None: + """A caller that composes the field itself pays for the field, and nothing else. + + Both the moved image and Transform.h5 are derived FROM the field -- a full-size resample and a + full-size rewrite of the same voxels. The tiled refinement reads the field, composes it with its + global pass and derives its own moved, so producing them for it is pure waste. + """ + moving = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + app = reg.ImpactRegKonfAIApp() + _stub_infer(app, moving, {"FireANTs_SyN": (2.0, 0.0, 0.0)}) + out = tmp_path / "Output" + app.register(["FireANTs_SyN"], [fixed], [moving], output=out, fields_only=True) + + case = out / "P000" + assert (case / "DVF.mha").is_file() + assert not (case / "Moved.mha").exists() + assert not (case / "Transform.h5").exists() From 3ac59c3e5d999ffc8c26617cf7d5ccc1467a3dba Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 00:28:42 +0200 Subject: [PATCH 15/39] feat(data): a reference that follows the case Resample's reference accepts '{case}': each case adopts the grid of its own entry in reference_group -- the registration idiom, where a moved image belongs on its field's grid. The reference-grid memo is per resolved entry, so a literal reference stays one lookup for the cohort. write_stream_cache_attribute gains the case name, and the landing folds (patching, transformer) and streamed writes (predictor) pass it: a per-case target grid resolves during planning and streaming, not only in __call__. --- konfai/data/patching.py | 14 +++--- konfai/data/transform.py | 94 ++++++++++++++++++++++++++++------------ konfai/predictor.py | 6 +-- konfai/transformer.py | 2 +- 4 files changed, 79 insertions(+), 37 deletions(-) diff --git a/konfai/data/patching.py b/konfai/data/patching.py index d1809900..6ad8b66e 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -107,7 +107,9 @@ def stream_region_source( cache_attribute: Attribute, ) -> list[slice]: ... - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: ... + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: ... def stream_region( self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute @@ -216,7 +218,9 @@ def stream_region( del context return self(name, tensor, cache_attribute) - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: """An augmentation draws a copy of the case rather than restating its geometry: nothing to record.""" def stream_shape(self, shape: list[int]) -> list[int]: @@ -1825,7 +1829,7 @@ def _plan_read_stage( # ORIENTATION / CROP / REGRID: the stage's own remap, on the state the stages before it left. pull = _RemapPull(stage.stream_region_source, list(shape), Attribute(evolved), self.name) out = self._stage_out_shape(stage, shape, Attribute(evolved)) - stage.write_stream_cache_attribute(evolved, list(shape)) + stage.write_stream_cache_attribute(evolved, list(shape), self.name) return _ReadStagePlan(loc.kind, tuple(shape), tuple(out), pull) def _stage_out_shape(self, stage: Stage, shape: list[int], attribute: Attribute) -> list[int]: @@ -1849,7 +1853,7 @@ def _fold_case_state(self, stage: Stage, shape: list[int], attribute: Attribute) """ out = self._stage_out_shape(stage, shape, attribute) if isinstance(stage, Transform): - stage.write_stream_cache_attribute(attribute, list(shape)) + stage.write_stream_cache_attribute(attribute, list(shape), self.name) return out @staticmethod @@ -2961,7 +2965,7 @@ def _replay_streamed_region( crop = [slice(t.start - s.start, t.stop - s.start) for t, s in zip(target, source, strict=False)] tensor = tensor[(*[slice(None)] * lead, *crop)] if case_attribute is not None: - stage.write_stream_cache_attribute(case_attribute, list(plan.in_shape)) + stage.write_stream_cache_attribute(case_attribute, list(plan.in_shape), self.name) self._check_region_geometry_reaches_the_case(stage, scoped, cache_attribute) return tensor, cache_attribute, keys_before diff --git a/konfai/data/transform.py b/konfai/data/transform.py index d25bef76..b8db0c78 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -277,7 +277,9 @@ def stream_abort(self, name: str) -> None: region sink or buffer does not outlive the case. The base holds nothing. """ - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: """Record the geometry a whole-volume ``__call__`` would, given the FULL source shape. Called once per case, on the persistent attribute, for the stage that owns a streamed region. @@ -286,6 +288,9 @@ def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatia ``__call__`` is handed while streaming: it writes the case-level answer here instead, and the patch-local one it wrote on the way is dropped rather than persisted. The base is a no-op -- a transform that leaves geometry alone has nothing to record. + + ``name`` is the case the fold walks — what a per-case answer (a ``Resample`` whose + reference follows the case) resolves against; a stage whose answer is case-blind ignores it. """ def stream_region( @@ -1054,6 +1059,12 @@ class _ReferenceGrid(_TargetGrid): silently, because a transposed grid resamples perfectly well onto the wrong place. Naming an image cannot make it: 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. + + An entry containing ``{case}`` FOLLOWS THE CASE: each case adopts the grid of its own entry -- + ``reference: '{case}', reference_group: DVF`` puts every moved image on its own field's grid, + which is the registration idiom (a displacement field is defined ON the fixed grid). The + literal spelling stays one lookup for the whole cohort; a per-case one is one per case, + headers only either way. """ needs = frozenset(_GEOMETRY_KEYS) @@ -1068,7 +1079,7 @@ def __init__(self, entry: str, group: str | None, dataset: str | None) -> None: filename, _flag, file_format = split_path_spec(str(dataset), default_format="mha") self.dataset = Dataset(Path(filename), file_format) self.roots: list[Dataset] = [] - self._grid: Grid | None = None + self._grids: dict[str, Grid] = {} def set_datasets(self, datasets: list[Dataset]) -> None: self.roots = list(datasets) @@ -1089,49 +1100,65 @@ def _group_in(self, dataset: Dataset) -> str: "Name it: Resample: {reference: " + self.entry + ", reference_group: }.", ) - def grid(self) -> Grid: - """The reference's grid, read from its header once. + def _entry_for(self, name: str) -> str: + """The entry to adopt for ``name`` — literal, or the case's own when it says ``{case}``.""" + if "{case}" not in self.entry: + return self.entry + if not name: + raise TransformError( + f"'Resample' has a per-case reference ('{self.entry}') and no case to resolve it for.", + "A per-case reference adopts, for each case, the grid of that case's own entry in" + " reference_group; it has no single grid to answer a caseless probe with.", + ) + return self.entry.replace("{case}", name) + + def grid(self, name: str = "") -> Grid: + """The reference's grid, read from its header once per distinct entry. - 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. + Headers only, and memoized by ENTRY: a literal reference is one lookup for the whole + cohort, a per-case one is one per case -- never the same answer bought again. """ - if self._grid is not None: - return self._grid + entry = self._entry_for(name) + cached = self._grids.get(entry) + if cached is not None: + return cached roots = self._roots() if not roots: raise TransformError( - f"'Resample' has no dataset to look reference '{self.entry}' up in.", + f"'Resample' has no dataset to look reference '{entry}' up in.", "Give the stage a root of its own -- Resample: {reference: " - + self.entry + + entry + ", reference_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) - self._grid = Grid.of([int(extent) for extent in shape[1:]], attribute, f"reference '{self.entry}'") - return self._grid + if dataset.is_dataset_exist(group, entry): + shape, attribute = dataset.get_infos(group, entry) + grid = Grid.of([int(extent) for extent in shape[1:]], attribute, f"reference '{entry}'") + self._grids[entry] = grid + return grid raise TransformError( - f"'Resample' cannot find reference '{self.entry}'" + f"'Resample' cannot find reference '{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.", + "Check the entry name and its group. A literal reference is looked up by entry -- one" + " grid serves the whole cohort; a '{case}' reference expects every case to have its own" + " entry in that group.", ) def of(self, source: Grid, name: str) -> Grid: - grid = self.grid() + grid = self.grid(name) if grid.rank != source.rank: where = f"case '{name}'" if name else "the case" raise TransformError( f"'Resample' cannot resample {where}, which has {source.rank} spatial axis/axes," - f" onto reference '{self.entry}', which has {grid.rank}." + f" onto reference '{self._entry_for(name)}', which has {grid.rank}." ) return grid def describe(self) -> str: - return f"reference '{self.entry}'" + return f"reference '{self.entry}'" + (" (per case)" if "{case}" in self.entry else "") class Resample(TransformInverse): @@ -1145,6 +1172,9 @@ class Resample(TransformInverse): - ``spacing``: the same field of view at another density. A component left at ``0`` keeps its axis. - ``shape``: the same field of view at a given count. A component left at ``0`` keeps its axis. - ``reference``: the grid of a stored image, adopted whole — extent, spacing, origin, direction. + ``'{case}'`` in the entry follows the case: each case adopts the grid of its OWN entry in + ``reference_group`` — ``reference: '{case}', reference_group: DVF`` lands every moved image + on its own field's grid, which is where a displacement field is defined. **What map to write it through** — any of, composed in this order: @@ -1557,7 +1587,7 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) # equality between the two paths is then a property of the code, not a claim about it. whole = tuple(slice(0, extent) for extent in target.size_zyx) result = self._sample(name, tensor, whole, [0] * source.rank) - self.write_stream_cache_attribute(cache_attribute, shape) + self.write_stream_cache_attribute(cache_attribute, shape, name) return result def _sample( @@ -1589,15 +1619,17 @@ def _mode(self, tensor: torch.Tensor) -> str: declared = self.interpolation or ("nearest" if tensor.dtype == torch.uint8 else "linear") return "nearest" if declared == "nearest" else "linear" - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: """Push the target grid over the source's, so the case now IS the grid it was written on. Pushed and not replaced: the source geometry stays underneath for :meth:`inverse` to pop back to, which is the whole of the stack this and ``_inverse_geometry`` share. """ shape = [int(extent) for extent in source_spatial_shape] - source, missing = Grid.from_header(shape, cache_attribute, "the case") - target = self._target.of(source, "") + source, missing = Grid.from_header(shape, cache_attribute, f"case '{name}'" if name else "the case") + target = self._target.of(source, name) written = { "Spacing": target.spacing_xyz, "Origin": target.origin_xyz, @@ -2871,9 +2903,12 @@ def stream_region_source( ) return source_slices - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: # Nothing to state for a case this cannot reorient: no geometry, or a direction that is not # 3-D. Its __call__ fails loudly before reaching here; a landing fold must not fail for it. + del name if not Grid.readable(cache_attribute) or cache_attribute.get_np_array("Direction").size != 9: return initial_matrix = cache_attribute.get_tensor("Direction").reshape(3, 3).to(torch.double) @@ -2969,7 +3004,7 @@ def _reorient(self, tensor: torch.Tensor, reorientation: torch.Tensor) -> torch. def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: # Read the source geometry before recording the canonical one over it: the attribute stacks. reorientation = self._reorientation(cache_attribute) - self.write_stream_cache_attribute(cache_attribute, list(tensor.shape[1:])) + self.write_stream_cache_attribute(cache_attribute, list(tensor.shape[1:]), name) return self._reorient(tensor, reorientation) def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: @@ -3427,7 +3462,10 @@ def stream_region_source( for target, (start, _) in zip(target_slices, box, strict=False) ] - def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + def write_stream_cache_attribute( + self, cache_attribute: Attribute, source_spatial_shape: list[int], name: str = "" + ) -> None: + del name if "box" not in cache_attribute: return if not {"Origin", "Spacing", "Direction"} <= set(cache_attribute.keys()): @@ -3478,7 +3516,7 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) if "box" not in cache_attribute: return tensor box = self._parse_box(cache_attribute["box"]) - self.write_stream_cache_attribute(cache_attribute, list(tensor.shape[1:])) + self.write_stream_cache_attribute(cache_attribute, list(tensor.shape[1:]), name) # The box carries the FAR margin, so the stop it crops at is the one the extent in hand decides. for i, ((_, b), s) in enumerate(zip(box, tensor.shape[1:], strict=False)): box[i][1] = s - b diff --git a/konfai/predictor.py b/konfai/predictor.py index b53e8f82..bd9af934 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -977,7 +977,7 @@ def _make_pipe_state( else: pull_fns.append(_RemapPull(transform.stream_region_source, shape, snapshot, name)) out = transform.transform_shape(self.group_src, name, list(shape), Attribute(walking)) - transform.write_stream_cache_attribute(walking, shape) + transform.write_stream_cache_attribute(walking, shape, name) shapes.append([int(extent) for extent in out]) elif locality.kind in _REGION_KINDS: if stage.inverted: @@ -1056,14 +1056,14 @@ def _apply_pipe_stage( remapper.inverse_stream_cache_attribute(attribute, in_shape) else: result = stage.transform.stream_region(name, block, context, Attribute(attribute)) - stage.transform.write_stream_cache_attribute(attribute, in_shape) + stage.transform.write_stream_cache_attribute(attribute, in_shape, name) return result if kind is LocalityKind.ORIENTATION and not stage.inverted: # A forward orientation writes the case origin/direction from the extent it is handed; run # the tensor action on a throwaway scope so it does not record the SLAB's extent, then write # the case geometry from the full ``in_shape`` (its documented contract) -- as REGRID does. result = stage(name, block, Attribute(attribute)) - cast(TransformInverse, stage.transform).write_stream_cache_attribute(attribute, in_shape) + cast(TransformInverse, stage.transform).write_stream_cache_attribute(attribute, in_shape, name) return result result = stage(name, block, attribute) if kind is LocalityKind.HALO: diff --git a/konfai/transformer.py b/konfai/transformer.py index 408264cc..55f86049 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -640,7 +640,7 @@ def _plan_notes(self) -> list[str]: int(extent) for extent in stage.transform_shape(manager.group_src, manager.name, source, attributes) ] - stage.write_stream_cache_attribute(attributes, source) + stage.write_stream_cache_attribute(attributes, source, manager.name) return notes def setup(self, world_size: int): From 5df9455c4236eb33015910c5fc3ed1e479728b86 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 00:28:57 +0200 Subject: [PATCH 16/39] feat(data): Std reduction and Magnitude transform Std folds cases element-wise with Welford running moments: incremental and voxel-local, so the peak is two accumulators plus the member being read; unbiased to match torch.std; zeros for a single case. Magnitude is the vector norm over the CHANNEL axis, POINTWISE, so the magnitude of a displacement field read as a case streams slab by slab. Norm keeps the trailing-axis stacked layout. --- konfai/data/reduction.py | 40 ++++++++++++++++++++++++++++++++++++++++ konfai/data/transform.py | 19 +++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/konfai/data/reduction.py b/konfai/data/reduction.py index bbb85fa9..db944fa9 100644 --- a/konfai/data/reduction.py +++ b/konfai/data/reduction.py @@ -154,6 +154,46 @@ def finalize(self) -> torch.Tensor: return result +class Std(Reduction): + """The element-wise standard deviation across cases — the ensemble-spread map. + + Welford's running moments, so the peak is two float32 accumulators plus the case being read, + whatever N is. Unbiased (N-1), matching ``torch.std``; a single case has no spread and + finalizes to zeros. + """ + + voxel_local = True + incremental = True + + def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: + self.start() + for tensor in tensors: + self.accumulate(tensor) + return self.finalize() + + def start(self) -> None: + self._count = 0 + self._mean: torch.Tensor | None = None + self._m2: torch.Tensor | None = None + + def accumulate(self, tensor: torch.Tensor) -> None: + value = tensor.float() + self._count += 1 + if self._mean is None or self._m2 is None: + self._mean, self._m2 = value.clone(), torch.zeros_like(value) + return + delta = value - self._mean + self._mean.add_(delta / self._count) + self._m2.addcmul_(delta, value - self._mean) + + def finalize(self) -> torch.Tensor: + if self._mean is None or self._m2 is None: + raise ReductionError("Std.finalize() with no case accumulated.", "Accumulate at least one case.") + result = torch.zeros_like(self._mean) if self._count < 2 else (self._m2 / (self._count - 1)).sqrt_() + self._mean, self._m2, self._count = None, None, 0 + return result + + class Median(Reduction): """The element-wise median across cases. diff --git a/konfai/data/transform.py b/konfai/data/transform.py index b8db0c78..f20564e3 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -3284,6 +3284,25 @@ def stream_abort(self, name: str) -> None: sink.abort() +class Magnitude(Transform): + """Vector magnitude over the CHANNEL axis: ``[C, ...]`` becomes ``[1, ...]``. + + :class:`Norm`'s channel-first sibling. ``Norm`` folds the trailing axis of a stacked ensemble + and is whole-volume by construction (a rank change past the streamed write); a stored vector + volume — a displacement field read as a case — is channel-first, and its magnitude at a voxel + reads that voxel alone: POINTWISE, so it streams. + """ + + def __init__(self) -> None: + super().__init__() + + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + return PatchLocality(LocalityKind.POINTWISE) + + def __call__(self, name: str, tensors: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: + return torch.linalg.norm(tensors.float(), dim=0, keepdim=True) + + class Norm(Transform): """Vector magnitude over the trailing component axis. From c6ca8e4d72fd8a31958c4db7e509dc15085a152f Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 00:29:08 +0200 Subject: [PATCH 17/39] feat(api): the workflows as Python callables konfai.api, re-exported lazily at top level: transform, plan_transform, evaluate, predict, train. A chain is a list of live stage objects or the equivalent mapping; the tree goes through the binder unchanged, and the resolved YAML still lands in the workspace as the run's record. Two halves make the object spelling possible without a second grammar: configure_workflow_environment materializes a {root: {...}} dict into a config file, so every workflow entry point takes the tree as well as a path; and record_given_arguments (applied through __init_subclass__ on Transform, DataAugmentation and Criterion) stores on each instance the kwargs the caller passed -- the binder's mirror. The outermost constructor records; a delegating super().__init__ keeps the caller's spelling; *args marks the instance unrecordable and it is refused by name. The contract differs from the CLI: a designed refusal raises KonfAIError instead of exiting, results come back structured (the outputs.json destinations, the parsed Metric_*.json), the KONFAI_* environment is restored around every call, and one workflow runs at a time per process -- a second concurrent call is refused with the remedy. tests/unit/test_api.py pins the recording, the serialization, byte-identity between the object and tree spellings of one run, and the contract. --- konfai/__init__.py | 25 ++ konfai/api.py | 439 ++++++++++++++++++++++++++++++++++++ konfai/data/augmentation.py | 8 +- konfai/data/transform.py | 8 +- konfai/evaluator.py | 4 +- konfai/metric/measure.py | 8 +- konfai/predictor.py | 4 +- konfai/trainer.py | 4 +- konfai/transformer.py | 19 +- konfai/utils/config.py | 50 ++++ konfai/utils/runtime.py | 38 +++- tests/unit/test_api.py | 234 +++++++++++++++++++ 12 files changed, 825 insertions(+), 16 deletions(-) create mode 100644 konfai/api.py create mode 100644 tests/unit/test_api.py diff --git a/konfai/__init__.py b/konfai/__init__.py index 68978efc..2692c126 100755 --- a/konfai/__init__.py +++ b/konfai/__init__.py @@ -387,3 +387,28 @@ def assert_konfai_install() -> None: lines.append(f" - {p}: {e}") raise KonfAIPackagesError("\n".join(lines)) + + +#: The Python workflow API (:mod:`konfai.api`), re-exported lazily: ``konfai.transform(...)`` +#: works, and ``import konfai`` stays light -- torch and the engines load on first use only. +_API_EXPORTS = ( + "transform", + "plan_transform", + "evaluate", + "predict", + "train", + "TransformResult", + "EvaluationResult", +) + + +def __getattr__(name: str): + if name in _API_EXPORTS: + from konfai import api + + return getattr(api, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(list(globals()) + list(_API_EXPORTS)) diff --git a/konfai/api.py b/konfai/api.py new file mode 100644 index 00000000..12719eff --- /dev/null +++ b/konfai/api.py @@ -0,0 +1,439 @@ +# 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 + +"""KonfAI in Python — the workflows as callables. + +The four CLI commands, callable: :func:`transform` (with :func:`plan_transform`, its dry-run +twin), :func:`evaluate`, :func:`predict` and :func:`train`. One engine, two spellings: everything +here builds the same config tree the YAML file would hold and hands it to the same binder — a +chain is a list of live stage objects (their constructor arguments are recorded as given, see +``record_given_arguments``), or the equivalent mapping, or a whole tree loaded from an existing +YAML and modified in place. + +The contract, and how it differs from the CLI: + +- **A designed refusal raises** ``KonfAIError`` — the message and the remedy are the exception; + the caller decides. Only the CLI catches and exits. +- **Results come back structured**: what a run wrote and where, read from the run's own record + (``outputs.json``, ``Metric_*.json``) instead of leaving the caller to fish files out. +- **The process is left as found**: the ``KONFAI_*`` environment is restored around every call, + and one workflow runs at a time per process — a second concurrent call is refused with the + remedy (subprocesses), not allowed to corrupt the first. +- **The record remains.** Every call materializes the resolved YAML in the run's workspace: a + notebook run is promoted to a versioned experiment by copying ``result.config`` — nothing to + rewrite. +""" + +import importlib +import json +import os +import threading +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from konfai.utils.errors import ConfigError + +if TYPE_CHECKING: + from konfai.transformer import TransformPlan + +#: Where a bare stage name resolves, per stage family — the same rule the YAML loader applies. +_STAGE_MODULES = ("konfai.data.transform", "konfai.data.augmentation") +_CRITERION_MODULES = ("konfai.metric.measure",) + +_ACTIVE = threading.Lock() + + +@contextmanager +def _one_workflow_at_a_time(ranks: int) -> Iterator[None]: + """Serialize workflows within the process and leave the environment as found. + + The engine keys its state on process-wide ``KONFAI_*`` variables, so two in-process runs would + corrupt each other — refused with the remedy rather than allowed. ``ranks`` is exported as + ``KONFAI_LOCAL_RANKS`` for build-time budget sizing, exactly as the CLI launcher does. + """ + if not _ACTIVE.acquire(blocking=False): + raise ConfigError( + "A KonfAI workflow is already running in this process.", + "Wait for it to return, or run concurrent workflows in separate processes: the engine" + " keys its state on process-wide KONFAI_* variables, so two in-process runs would" + " corrupt each other.", + ) + saved = {key: value for key, value in os.environ.items() if key.startswith("KONFAI")} + os.environ["KONFAI_LOCAL_RANKS"] = str(max(1, ranks)) + try: + yield + finally: + for key in [key for key in os.environ if key.startswith("KONFAI")]: + if key not in saved: + del os.environ[key] + os.environ.update(saved) + _ACTIVE.release() + + +def _yaml_safe(value: object, where: str) -> object: + """``value`` as the config file could hold it — or a refusal that names the argument.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, Path): + return str(value) + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Mapping): + return {str(key): _yaml_safe(entry, f"{where}.{key}") for key, entry in value.items()} + if isinstance(value, (list, tuple)): + return [_yaml_safe(entry, f"{where}[{index}]") for index, entry in enumerate(value)] + raise ConfigError( + f"'{where}' is a {type(value).__name__}, which a config tree cannot hold.", + "A stage's constructor arguments must be YAML-spellable (numbers, strings, paths, lists," + " mappings) -- the same rule the YAML file obeys.", + ) + + +def _classpath(stage: object, default_modules: tuple[str, ...]) -> str: + """The name the config tree references ``stage`` by — bare for a shipped class, qualified else.""" + cls = type(stage) + return cls.__name__ if cls.__module__ in default_modules else f"{cls.__module__}:{cls.__name__}" + + +def _qualified_spelling(stage: object, name: str, default_modules: tuple[str, ...]) -> str: + """A repeated bare name's second spelling — module-qualified, resolved as the binder resolves it.""" + if not isinstance(stage, Mapping): + return f"{type(stage).__module__}:{type(stage).__name__}" + for module in default_modules: + if hasattr(importlib.import_module(module), name): + return f"{module}:{name}" + raise ConfigError( + f"'{name}' appears twice and resolves in no default module, so its second spelling is unknown.", + "Spell the second occurrence with its module: {'my_module:" + name + "': {...}}.", + ) + + +def _stage_entry(stage: object, default_modules: tuple[str, ...], where: str) -> tuple[str, object]: + """One chain entry: ``(classpath, kwargs-subtree)`` from a live object or a mapping.""" + if isinstance(stage, Mapping): + if len(stage) != 1: + raise ConfigError( + f"'{where}' is a mapping of {len(stage)} stages; a chain entry holds exactly one.", + "Spell each stage as its own entry: [{'Clip': {...}}, {'Write': {...}}] -- or" + " instantiate the classes and pass the objects.", + ) + ((name, kwargs),) = stage.items() + return str(name), _yaml_safe(kwargs, f"{where}.{name}") + given = getattr(stage, "_konfai_given", None) + if given is None: + raise ConfigError( + f"'{where}' ({type(stage).__name__}) records no constructor arguments.", + "A chain stage is a Transform, DataAugmentation or Criterion subclass instance -- their" + " bases record what the constructor was given -- or a one-entry mapping" + " {'Name': {...kwargs...}}.", + ) + name = _classpath(stage, default_modules) + return name, {key: _yaml_safe(value, f"{where}.{name}.{key}") for key, value in given.items()} + + +def _chain_tree(stages: object, default_modules: tuple[str, ...], where: str) -> dict[str, object]: + """A chain — a sequence of stages — as the mapping the config tree holds, in order. + + The tree is a mapping, so two stages of the same class need distinct spellings: the second + occurrence is written module-qualified (which resolves to the same class); a third has no + spelling left and is refused — the YAML file has the same limit. + """ + if isinstance(stages, Mapping): # a chain already spelled as its tree + return {str(key): _yaml_safe(value, f"{where}.{key}") for key, value in stages.items()} + tree: dict[str, object] = {} + for index, stage in enumerate(_stage_sequence(stages, where)): + name, kwargs = _stage_entry(stage, default_modules, f"{where}[{index}]") + if name in tree and ":" not in name: + name = _qualified_spelling(stage, name, default_modules) + if name in tree: + raise ConfigError( + f"'{where}' holds three stages spelled '{name.split(':')[-1]}'; the tree is a" + " mapping and has two spellings (bare and module-qualified), not three.", + "Split the chain around a Save/Write boundary, or subclass the stage under a distinct name.", + ) + tree[name] = kwargs + return tree + + +def _stage_sequence(stages: object, where: str) -> Sequence[object]: + if isinstance(stages, Sequence) and not isinstance(stages, (str, bytes)): + return stages + raise ConfigError( + f"'{where}' is a {type(stages).__name__}; a chain is a sequence of stages.", + "Pass the stages in application order: [Clip(min_value=0), Write(dataset='./Out:mha')].", + ) + + +def _dataset_filenames(datasets: str | Path | Sequence[str | Path]) -> list[str]: + entries = [datasets] if isinstance(datasets, (str, Path)) else list(datasets) + if not entries: + raise ConfigError( + "No dataset was given.", + "Name at least one root, as the YAML would: './Dataset:mha' (path, then format).", + ) + return [str(entry) for entry in entries] + + +# ------------------------------------------------------------------------------------- TRANSFORM + + +@dataclass(frozen=True) +class TransformResult: + """What a TRANSFORM run produced, in the run's own terms.""" + + #: The run directory (``Transforms/``): logs, ``plan.txt``, the resolved config, + #: ``outputs.json`` — never the deliverable, which each ``Write`` placed in the caller's tree. + workspace: Path + #: Every chain's terminal ``Write``: ``{group_src, group_dest, dataset, group, format}``. + outputs: list[dict[str, str]] + #: The resolved config the run kept — copy this file to version the experiment. + config: Path + + +def _transform_tree( + name: str, + datasets: str | Path | Sequence[str | Path], + chains: Mapping[str, Mapping[str, object]], + memory_budget: str | int | None, + on_fallback: str, + manual_seed: int, + dataset_options: Mapping[str, object] | None, +) -> dict: + groups_src: dict[str, object] = {} + for group_src, destinations in chains.items(): + groups_dest = { + str(group_dest): {"transforms": _chain_tree(stages, _STAGE_MODULES, f"chains.{group_src}.{group_dest}")} + for group_dest, stages in destinations.items() + } + groups_src[str(group_src)] = {"groups_dest": groups_dest} + dataset_tree: dict[str, object] = {"dataset_filenames": _dataset_filenames(datasets), "groups_src": groups_src} + if memory_budget is not None: + dataset_tree["memory_budget"] = memory_budget + dataset_tree.update(dict(dataset_options or {})) + return { + "Transformer": { + "name": name, + "on_fallback": on_fallback, + "manual_seed": manual_seed, + "Dataset": dataset_tree, + } + } + + +def transform( + name: str, + datasets: str | Path | Sequence[str | Path], + chains: Mapping[str, Mapping[str, object]], + *, + memory_budget: str | int | None = None, + on_fallback: str = "warn", + manual_seed: int = 0, + dataset_options: Mapping[str, object] | None = None, + gpu: Sequence[int] | None = None, + cpu: int = 1, + quiet: bool = False, + overwrite: bool = False, + transforms_dir: Path | str = Path("./Transforms"), +) -> TransformResult: + """Run a TRANSFORM workflow: read a dataset, apply each chain, ``Write`` the results. + + ``chains`` maps ``group_src -> group_dest -> chain``, where a chain is a list of stage objects + (``[Resample(...), Write(dataset='./Out:mha')]``), of one-entry mappings, or the equivalent + mapping tree. Every chain ends in a ``Write`` — the same rule, message and plan as the CLI, + which is this function with a YAML file. GPU is opt-in (``gpu=[0]``), as it is on the CLI. + """ + tree = _transform_tree(name, datasets, chains, memory_budget, on_fallback, manual_seed, dataset_options) + from konfai.transformer import build_transform + from konfai.utils.runtime import execute_distributed_object + + with _one_workflow_at_a_time(len(gpu or []) or cpu): + workflow = build_transform(transform_file=tree, transforms_dir=transforms_dir) + config_name = Path(os.environ["KONFAI_config_file"]).name + execute_distributed_object(workflow, gpu=list(gpu or []), cpu=cpu, overwrite=overwrite, quiet=quiet) + workspace = Path(workflow.transform_path) # type: ignore[attr-defined] + outputs = json.loads((workspace / "outputs.json").read_text(encoding="utf-8")) + return TransformResult(workspace=workspace, outputs=outputs, config=workspace / config_name) + + +def plan_transform( + name: str, + datasets: str | Path | Sequence[str | Path], + chains: Mapping[str, Mapping[str, object]], + *, + memory_budget: str | int | None = None, + on_fallback: str = "warn", + manual_seed: int = 0, + dataset_options: Mapping[str, object] | None = None, + cpu: int = 1, + overwrite: bool = False, + transforms_dir: Path | str = Path("./Transforms"), +) -> "TransformPlan": + """:func:`transform`'s dry-run twin: build, plan, print, return the plan — the run never starts. + + Same arguments, same verdicts: the returned ``TransformPlan`` is the run's own routing + (STREAM/LOAD/WHOLE-VOLUME/SKIP/REDUCE), measured, not estimated. + """ + tree = _transform_tree(name, datasets, chains, memory_budget, on_fallback, manual_seed, dataset_options) + from konfai.transformer import plan_transform as _plan_transform + + with _one_workflow_at_a_time(cpu): + return _plan_transform(transform_file=tree, transforms_dir=transforms_dir, cpu=cpu, overwrite=overwrite) + + +# ------------------------------------------------------------------------------------ EVALUATION + + +@dataclass(frozen=True) +class EvaluationResult: + """What an EVALUATION run measured, parsed from its own record.""" + + #: The run directory (``Evaluations/``), holding the ``Metric_*.json`` files. + workspace: Path + #: The parsed metric reports, keyed by split (``TRAIN``, ``VALIDATION``) when present. + metrics: dict[str, Any] + + +def evaluate( + name: str, + datasets: str | Path | Sequence[str | Path], + metrics: Mapping[str, Mapping[str, object]], + *, + transforms: Mapping[str, object] | None = None, + dataset_options: Mapping[str, object] | None = None, + gpu: Sequence[int] | None = None, + cpu: int = 1, + quiet: bool = False, + overwrite: bool = False, + evaluations_dir: Path | str = Path("./Evaluations"), +) -> EvaluationResult: + """Run an EVALUATION workflow and return the measured metrics. + + ``metrics`` maps ``output_group -> target_group -> criteria``, where criteria is a list of + :class:`~konfai.metric.measure.Criterion` instances (``[MAE(), Dice(labels=[1, 2])]``) or + one-entry mappings. ``transforms`` optionally names a pre-metric chain per group (a cast, a + clip) — the groups themselves are derived from ``metrics``, declared once, not twice. + """ + # A composite target ("Seg;Mask") is one metric key over several dataset groups: split it here, + # as the Evaluator does, so each named group is actually loaded. + groups = sorted( + {str(group) for group in metrics} + | {part for targets in metrics.values() for target in targets for part in str(target).split(";")} + ) + groups_src: dict[str, object] = {} + for group in groups: + destination: dict[str, object] = {"is_input": True} + if transforms is not None and group in transforms: + destination["transforms"] = _chain_tree(transforms[group], _STAGE_MODULES, f"transforms.{group}") + groups_src[group] = {"groups_dest": {group: destination}} + metrics_tree = { + str(output): { + "targets_criterions": { + str(target): { + "criterions_loader": _chain_tree(criteria, _CRITERION_MODULES, f"metrics.{output}.{target}") + } + for target, criteria in targets.items() + } + } + for output, targets in metrics.items() + } + dataset_tree: dict[str, object] = {"dataset_filenames": _dataset_filenames(datasets), "groups_src": groups_src} + dataset_tree.update(dict(dataset_options or {})) + tree = {"Evaluator": {"train_name": name, "metrics": metrics_tree, "Dataset": dataset_tree}} + + from konfai.evaluator import build_evaluate + from konfai.utils.runtime import execute_distributed_object + + with _one_workflow_at_a_time(len(gpu or []) or cpu): + workflow = build_evaluate(evaluations_file=tree, evaluations_dir=evaluations_dir) + execute_distributed_object(workflow, gpu=list(gpu or []), cpu=cpu, overwrite=overwrite, quiet=quiet) + workspace = Path(os.environ["KONFAI_EVALUATIONS_DIRECTORY"]) / name + reports = { + split: json.loads(report.read_text(encoding="utf-8")) + for split in ("TRAIN", "VALIDATION") + if (report := workspace / f"Metric_{split}.json").is_file() + } + return EvaluationResult(workspace=workspace, metrics=reports) + + +# ------------------------------------------------------------------------- PREDICTION / TRAINING + + +def predict( + models: Sequence[Path | str], + config: Mapping[str, object] | Path | str, + *, + gpu: Sequence[int] | None = None, + cpu: int = 1, + quiet: bool = False, + overwrite: bool = False, + predictions_dir: Path | str = Path("./Predictions"), +) -> Path: + """Run a PREDICTION workflow; return its workspace (``Predictions/``). + + ``config`` is a ``Prediction.yml`` path or the same tree as a dict — a prediction's substance + (checkpoints, patching, TTA, ensembling) is wiring, and the tree is its honest spelling. + """ + from konfai.predictor import build_predict + from konfai.utils.runtime import execute_distributed_object + + with _one_workflow_at_a_time(len(gpu or []) or cpu): + workflow = build_predict( + models=[Path(model) for model in models], + prediction_file=config if isinstance(config, (Path, str)) else dict(config), + predictions_dir=predictions_dir, + ) + execute_distributed_object(workflow, gpu=list(gpu or []), cpu=cpu, overwrite=overwrite, quiet=quiet) + return Path(os.environ["KONFAI_PREDICTIONS_DIRECTORY"]) / workflow.name + + +def train( + config: Mapping[str, object] | Path | str, + *, + resume: bool = False, + model: Path | str | None = None, + lr: float | None = None, + gpu: Sequence[int] | None = None, + cpu: int | None = None, + quiet: bool = False, + overwrite: bool = False, + checkpoints_dir: Path | str = Path("./Checkpoints"), + statistics_dir: Path | str = Path("./Statistics"), +) -> Path: + """Run a TRAIN (or RESUME) workflow; return its checkpoint workspace. + + ``config`` is a ``Config.yml`` path or the same tree as a dict. The Python idiom for a sweep + is the tree: load the YAML once, change the keys under study, call this — the resolved config + each run keeps IS the record of what was tried. + """ + from konfai.trainer import build_train + from konfai.utils.runtime import State, execute_distributed_object + + with _one_workflow_at_a_time(len(gpu or []) or (cpu or 1)): + workflow = build_train( + command=State.RESUME if resume else State.TRAIN, + model=model, + config=config if isinstance(config, (Path, str)) else dict(config), + checkpoints_dir=checkpoints_dir, + statistics_dir=statistics_dir, + lr=lr, + ) + execute_distributed_object(workflow, gpu=list(gpu or []), cpu=cpu, overwrite=overwrite, quiet=quiet) + return Path(os.environ["KONFAI_CHECKPOINTS_DIRECTORY"]) / workflow.name diff --git a/konfai/data/augmentation.py b/konfai/data/augmentation.py index eb2714f9..239f550b 100755 --- a/konfai/data/augmentation.py +++ b/konfai/data/augmentation.py @@ -33,7 +33,7 @@ from konfai import konfai_root from konfai.data.transform import LocalityKind, PatchLocality -from konfai.utils.config import _escape_key_component, apply_config +from konfai.utils.config import _escape_key_component, apply_config, record_given_arguments from konfai.utils.dataset import Attribute, Dataset, data_to_image from konfai.utils.errors import AugmentationError from konfai.utils.runtime import NeedDevice @@ -213,6 +213,12 @@ def set_datasets(self, datasets: list[Dataset]) -> None: class DataAugmentation(NeedDevice, ABC): + def __init_subclass__(cls, **kwargs: object) -> None: + # A draw is a chain stage too: record its constructor arguments as given, so konfai.api can + # write the config tree back from live objects (see Transform.__init_subclass__). + super().__init_subclass__(**kwargs) + record_given_arguments(cls) + def __init__(self, groups: list[str] | None = None) -> None: self.who_index: dict[int, list[int]] = {} self.shape_index: dict[int, list[list[int]]] = {} diff --git a/konfai/data/transform.py b/konfai/data/transform.py index f20564e3..2693ca99 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -53,7 +53,7 @@ source_index, source_window, ) -from konfai.utils.config import _escape_key_component, apply_config +from konfai.utils.config import _escape_key_component, apply_config, record_given_arguments from konfai.utils.dataset import Attribute, Dataset, DataStream, data_to_image, image_to_data from konfai.utils.errors import TransformError from konfai.utils.ITK import _require_simpleitk, box_with_mask, crop_with_mask @@ -166,6 +166,12 @@ class Transform(NeedDevice, ABC): supports_dataloader_workers = True + def __init_subclass__(cls, **kwargs: object) -> None: + # Every stage records its constructor arguments as given, so konfai.api can write the + # config tree back from live objects -- the binder's mirror, declared once, on the base. + super().__init_subclass__(**kwargs) + record_given_arguments(cls) + def __init__(self) -> None: NeedDevice.__init__(self) self.datasets: list[Dataset] = [] diff --git a/konfai/evaluator.py b/konfai/evaluator.py index 973319a1..3f29e23b 100644 --- a/konfai/evaluator.py +++ b/konfai/evaluator.py @@ -648,7 +648,7 @@ def description(measure): def build_evaluate( - evaluations_file: Path | str = Path("./Evaluation.yml").resolve(), + evaluations_file: Path | str | dict = Path("./Evaluation.yml").resolve(), evaluations_dir: Path | str = Path("./Evaluations").resolve(), ) -> DistributedObject: """ @@ -683,7 +683,7 @@ def evaluate( cpu: int = 1, quiet: bool = False, tensorboard: bool = False, - evaluations_file: Path | str = Path("./Evaluation.yml").resolve(), + evaluations_file: Path | str | dict = Path("./Evaluation.yml").resolve(), evaluations_dir: Path | str = Path("./Evaluations").resolve(), ) -> DistributedObject: """ diff --git a/konfai/metric/measure.py b/konfai/metric/measure.py index 05a245f1..d86e7992 100644 --- a/konfai/metric/measure.py +++ b/konfai/metric/measure.py @@ -34,7 +34,7 @@ from konfai.data.patching import ModelPatch from konfai.network.blocks import LatentDistribution from konfai.network.network import ModelLoader, Network -from konfai.utils.config import apply_config +from konfai.utils.config import apply_config, record_given_arguments from konfai.utils.dataset import Attribute from konfai.utils.errors import MeasureError from konfai.utils.utils import get_module @@ -75,6 +75,12 @@ class Criterion(torch.nn.Module, ABC): # ``forward`` exactly may set it. reducible: bool = False + def __init_subclass__(cls, **kwargs: object) -> None: + # A metric is config-built like a stage: record its constructor arguments as given, so + # konfai.api can write the config tree back from live objects (see Transform). + super().__init_subclass__(**kwargs) + record_given_arguments(cls) + def __init__(self) -> None: super().__init__() diff --git a/konfai/predictor.py b/konfai/predictor.py index bd9af934..cdf931e9 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -2251,7 +2251,7 @@ def __repr__(self) -> str: def build_predict( models: list[Path], - prediction_file: Path | str = Path("./Prediction.yml").resolve(), + prediction_file: Path | str | dict = Path("./Prediction.yml").resolve(), predictions_dir: Path | str = Path("./Predictions").resolve(), ) -> DistributedObject: """ @@ -2291,7 +2291,7 @@ def predict( cpu: int = 1, quiet: bool = False, tensorboard: bool = False, - prediction_file: Path | str = Path("./Prediction.yml").resolve(), + prediction_file: Path | str | dict = Path("./Prediction.yml").resolve(), predictions_dir: Path | str = Path("./Predictions").resolve(), ) -> DistributedObject: """ diff --git a/konfai/trainer.py b/konfai/trainer.py index 53b0cf0c..cd4a954c 100644 --- a/konfai/trainer.py +++ b/konfai/trainer.py @@ -1093,7 +1093,7 @@ def _usable_vram_after_oom(self, device: int | None) -> float: def build_train( command: State = State.TRAIN, model: Path | str | None = None, - config: Path | str = Path("./Config.yml"), + config: Path | str | dict = Path("./Config.yml"), checkpoints_dir: Path | str = Path("./Checkpoints/"), statistics_dir: Path | str = Path("./Statistics/"), lr: float | None = None, @@ -1151,7 +1151,7 @@ def train( cpu: int | None = None, quiet: bool = False, tensorboard: bool = False, - config: Path | str = Path("./Config.yml"), + config: Path | str | dict = Path("./Config.yml"), checkpoints_dir: Path | str = Path("./Checkpoints/"), statistics_dir: Path | str = Path("./Statistics/"), lr: float | None = None, diff --git a/konfai/transformer.py b/konfai/transformer.py index 55f86049..eabd3a56 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -57,6 +57,7 @@ from konfai.utils.runtime import ( DistributedObject, State, + _materialized_config, configure_workflow_environment, run_distributed_app, ) @@ -988,7 +989,7 @@ def walk(node: object, grammar: dict[str, object], path: str) -> None: def build_transform( - transform_file: Path | str = Path("./Transform.yml").resolve(), + transform_file: Path | str | dict = Path("./Transform.yml").resolve(), transforms_dir: Path | str = Path("./Transforms").resolve(), ) -> DistributedObject: """Build and return the configured transform workflow without executing it. @@ -996,7 +997,13 @@ def build_transform( The returned object carries the plan: ``compute_plan()`` is the programmatic dry-run, and ``setup()`` prints and enforces it. This is the in-process surface — ``transform()`` stays ``None``-returning like every workflow entrypoint. + + ``transform_file`` may be the config TREE itself, as a dict, instead of a path: the Python + caller writes no YAML. Materialized here rather than in ``configure_workflow_environment`` + so the strict-grammar check below reads the same file every other reader will. """ + if isinstance(transform_file, dict): + transform_file = _materialized_config(transform_file, "Transformer") configure_workflow_environment( config_path=transform_file, root="Transformer", @@ -1011,7 +1018,7 @@ def build_transform( def plan_transform( overwrite: bool = False, cpu: int = 1, - transform_file: Path | str = Path("./Transform.yml").resolve(), + transform_file: Path | str | dict = Path("./Transform.yml").resolve(), transforms_dir: Path | str = Path("./Transforms").resolve(), **_ignored: object, ) -> TransformPlan: @@ -1033,9 +1040,13 @@ def transform( gpu: list[int] | None = None, cpu: int = 1, quiet: bool = False, - transform_file: Path | str = Path("./Transform.yml").resolve(), + transform_file: Path | str | dict = Path("./Transform.yml").resolve(), transforms_dir: Path | str = Path("./Transforms").resolve(), ) -> DistributedObject: - """Build and execute the configured transform workflow.""" + """Build and execute the configured transform workflow. + + ``transform_file`` accepts the config tree as a dict — the pure-Python spelling of the same + run; the resolved YAML still lands in the workspace as the run's record. + """ del overwrite, gpu, cpu, quiet return build_transform(transform_file=transform_file, transforms_dir=transforms_dir) diff --git a/konfai/utils/config.py b/konfai/utils/config.py index b8ca64b7..449b3fc8 100755 --- a/konfai/utils/config.py +++ b/konfai/utils/config.py @@ -17,6 +17,7 @@ """Configuration helpers that map YAML trees to KonfAI Python objects.""" import collections +import functools import inspect import logging import os @@ -667,3 +668,52 @@ def new_function(*args, **kwargs): return new_function return decorator + + +def record_given_arguments(cls: type) -> None: + """Make ``cls`` record, on each instance, the constructor arguments AS GIVEN — the binder's mirror. + + The binder builds an object from a config subtree; this makes the reverse spelling possible: an + object built in Python remembers what the caller said (``_konfai_given``), so :mod:`konfai.api` + can write a workflow tree from live objects with no second grammar — the recorded kwargs go + back through the binder, which stays the one place that validates and resolves defaults. + + Only the OUTERMOST constructor records: a subclass delegating to ``super().__init__`` keeps its + own spelling, which is what the caller wrote. Applied by the extension bases' + ``__init_subclass__``, so a subclass that defines no ``__init__`` inherits a recording one. An + ``__init__`` taking ``*args`` cannot be spelled as a config subtree; such an instance records + nothing and :mod:`konfai.api` refuses it by name. + """ + # A subclass with no __init__ of its own wraps the inherited one here: the extension bases are + # never passed through this function, so their raw constructors record nothing by themselves. + original = cls.__dict__.get("__init__") or cls.__init__ # type: ignore[misc] + if getattr(original, "_konfai_records", False): + return + signature = inspect.signature(original) + + @functools.wraps(original) + def recording(self: object, *args: object, **kwargs: object) -> None: + if not hasattr(self, "_konfai_given"): + arguments: dict[str, object] | None = {} + try: + bound = signature.bind(self, *args, **kwargs) + except TypeError: + arguments = None # the original call raises its own, better error just below + if arguments is not None: + for name, value in list(bound.arguments.items())[1:]: + kind = signature.parameters[name].kind + if kind is inspect.Parameter.VAR_POSITIONAL: + arguments = None + break + if kind is inspect.Parameter.VAR_KEYWORD: + arguments.update(dict(value)) # type: ignore[call-overload] + else: + arguments[name] = value + # None is recorded too: it marks the instance as spoken for, so a delegating + # super().__init__ cannot record the INNER spelling under the outer class's name -- + # kwargs the outer constructor does not accept. + self._konfai_given = arguments # type: ignore[attr-defined] + original(self, *args, **kwargs) + + recording._konfai_records = True # type: ignore[attr-defined] + cls.__init__ = recording # type: ignore[method-assign, misc] diff --git a/konfai/utils/runtime.py b/konfai/utils/runtime.py index 213949f8..3a77faf7 100644 --- a/konfai/utils/runtime.py +++ b/konfai/utils/runtime.py @@ -16,6 +16,7 @@ """Runtime state, logging, and distributed execution helpers for KonfAI.""" +import atexit import builtins import inspect import os @@ -25,6 +26,7 @@ import socket import subprocess # nosec B404 import sys +import tempfile import time from abc import ABC, abstractmethod from collections.abc import Callable @@ -180,9 +182,35 @@ def available_memory_bytes() -> tuple[int, str]: return host_available, "host available RAM" +def _materialized_config(tree: dict, root: str) -> Path: + """A config TREE written where a workflow expects a file — the Python front door. + + The caller hands the same tree the YAML file would hold — ``{root: {...}}``, the very kwargs + the binder feeds each ``__init__`` — and never touches YAML: it is written once here, under a + scratch directory of its own, and everything downstream (reflection binding, resolution + write-back, the workspace copy, resume) sees an ordinary config file. The workspace keeps the + resolved copy, as it does for every run. + """ + if list(tree) != [root]: + raise ConfigError( + f"A config tree for this workflow must hold exactly the '{root}' root" + f" (found: {sorted(str(key) for key in tree)}).", + f"Pass the same tree the YAML file would hold: {{'{root}': {{...}}}}.", + ) + from ruamel.yaml import YAML + + scratch = Path(tempfile.mkdtemp(prefix=f"konfai_{root.lower()}_")) + # The file must outlive the run (spawned ranks re-read it), not the process. + atexit.register(shutil.rmtree, scratch, ignore_errors=True) + path = scratch / f"{root}.yml" + with path.open("w", encoding="utf-8") as file: + YAML().dump(tree, file) + return path + + def configure_workflow_environment( *, - config_path: Path | str, + config_path: Path | str | dict, root: str, state: "State | str", path_env: dict[str, Path | str] | None = None, @@ -192,8 +220,10 @@ def configure_workflow_environment( Parameters ---------- - config_path : Path | str - YAML configuration file used by the workflow. + config_path : Path | str | dict + YAML configuration file used by the workflow — or the config TREE itself, as a dict, for + a Python caller that writes no YAML (see :func:`_materialized_config`). Every workflow + entry point accepts either, since they all pass through here. root : str Root configuration section, for example ``Trainer`` or ``Predictor``. state : State | str @@ -202,6 +232,8 @@ def configure_workflow_environment( Additional environment variables whose values should be normalized as absolute filesystem paths before export. """ + if isinstance(config_path, dict): + config_path = _materialized_config(config_path, root) os.environ["KONFAI_config_file"] = str(Path(config_path).resolve()) os.environ["KONFAI_ROOT"] = root os.environ["KONFAI_STATE"] = str(state) diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py new file mode 100644 index 00000000..c213b432 --- /dev/null +++ b/tests/unit/test_api.py @@ -0,0 +1,234 @@ +# 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 + +"""The Python front door (:mod:`konfai.api`): live objects and the YAML file are two spellings of +one engine. Pins the kwargs recording, the object->tree serialization, the run contract (raise, not +exit; the process env left as found; one workflow at a time), and byte-identity between the two +spellings of the same run.""" + +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +sitk = pytest.importorskip("SimpleITK") + +from konfai import api # noqa: E402 +from konfai.data.reduction import Std # noqa: E402 +from konfai.data.transform import Clip, Magnitude, Resample, Warp, Write # noqa: E402 +from konfai.metric.measure import Dice # noqa: E402 +from konfai.utils.errors import ConfigError, KonfAIError # noqa: E402 + +# --------------------------------------------------------------------------- recording and trees + + +def test_a_stage_records_the_arguments_as_given() -> None: + stage = Clip(min_value=-100.0, max_value=300.0) + assert stage._konfai_given == {"min_value": -100.0, "max_value": 300.0} + + +def test_a_criterion_records_too() -> None: + assert "labels" in Dice(labels=[1, 2])._konfai_given + + +def test_a_subclass_with_no_init_of_its_own_records_the_inherited_one() -> None: + """Accuracy inherits Criterion's constructor whole -- the recording must come with it.""" + from konfai.metric.measure import Accuracy + + assert Accuracy()._konfai_given == {} + + +def test_a_repeated_mapping_stage_is_qualified_by_resolution() -> None: + """The second occurrence of a bare mapping name gets the module the binder would resolve.""" + tree = api._chain_tree( + [{"Clip": {"min_value": 0.0}}, {"Clip": {"max_value": 1.0}}], + api._STAGE_MODULES, + "chains.CT.CT", + ) + assert list(tree) == ["Clip", "konfai.data.transform:Clip"] + + +def test_a_subclass_delegating_to_super_keeps_its_own_spelling() -> None: + """``Warp(field=...)`` expands into ``Resample`` arguments internally; the recorded spelling is + the caller's, so the tree references ``Warp`` with the caller's kwargs and rebinds identically.""" + stage = Warp(field="./DVF:omezarr", max_displacement=120.0) + assert stage._konfai_given == {"field": "./DVF:omezarr", "max_displacement": 120.0} + + +def test_the_chain_tree_is_the_yaml_subtree() -> None: + tree = api._chain_tree([Clip(min_value=0.0), Write(dataset="./Out:mha")], api._STAGE_MODULES, "chains.CT.CT") + assert tree == {"Clip": {"min_value": 0.0}, "Write": {"dataset": "./Out:mha"}} + + +def test_an_unrecordable_stage_is_refused_by_name() -> None: + class VarArgs(Clip): + def __init__(self, *bounds: float) -> None: + super().__init__(min_value=min(bounds)) + + with pytest.raises(ConfigError, match="VarArgs"): + api._chain_tree([VarArgs(1.0, 2.0)], api._STAGE_MODULES, "chains.CT.CT") + + +def test_a_non_spellable_argument_is_refused_by_name() -> None: + with pytest.raises(ConfigError, match="max_value"): + api._chain_tree([Clip(max_value=np.ones(2))], api._STAGE_MODULES, "chains.CT.CT") + + +# ------------------------------------------------------------------------------------ run contract + + +def _write_case(path: Path, values: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + sitk.WriteImage(sitk.GetImageFromArray(values), str(path)) + + +@pytest.fixture() +def cohort(tmp_path: Path) -> Path: + rng = np.random.default_rng(7) + for case in ("P000", "P001"): + _write_case(tmp_path / "Raw" / case / "CT.mha", rng.normal(0.0, 200.0, (6, 7, 8)).astype(np.float32)) + return tmp_path + + +def test_objects_and_yaml_are_two_spellings_of_one_run(cohort: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.chdir(cohort) + result = api.transform( + "BY_OBJECTS", + "./Raw:mha", + {"CT": {"CT": [Clip(min_value=-50.0, max_value=100.0), Write(dataset="./OutA:mha")]}}, + transforms_dir=cohort / "Transforms", + quiet=True, + ) + api.transform( + "BY_TREE", + "./Raw:mha", + {"CT": {"CT": {"Clip": {"min_value": -50.0, "max_value": 100.0}, "Write": {"dataset": "./OutB:mha"}}}}, + transforms_dir=cohort / "Transforms", + quiet=True, + ) + for case in ("P000", "P001"): + by_objects = (cohort / "OutA" / case / "CT.mha").read_bytes() + by_tree = (cohort / "OutB" / case / "CT.mha").read_bytes() + assert by_objects == by_tree + assert result.workspace == cohort / "Transforms" / "BY_OBJECTS" + assert result.outputs[0]["dataset"] == str(cohort / "OutA") + assert result.config.is_file() + + +def test_a_designed_refusal_raises_instead_of_exiting(cohort: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.chdir(cohort) + with pytest.raises(KonfAIError): + api.transform( + "NO_WRITE", + "./Raw:mha", + {"CT": {"CT": [Clip(min_value=0.0)]}}, + transforms_dir=cohort / "Transforms", + quiet=True, + ) + + +def test_the_environment_is_left_as_found(cohort: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.chdir(cohort) + for key in [key for key in os.environ if key.startswith("KONFAI")]: + monkeypatch.delenv(key) + api.transform( + "ENV", + "./Raw:mha", + {"CT": {"CT": [Write(dataset="./OutEnv:mha")]}}, + transforms_dir=cohort / "Transforms", + quiet=True, + ) + assert [key for key in os.environ if key.startswith("KONFAI")] == [] + + +def test_one_workflow_at_a_time_per_process(cohort: Path) -> None: + assert api._ACTIVE.acquire(blocking=False) + try: + with pytest.raises(ConfigError, match="already running"): + api.transform("BUSY", "./Raw:mha", {"CT": {"CT": [Write(dataset="./Out:mha")]}}) + finally: + api._ACTIVE.release() + + +def test_the_reference_follows_the_case(cohort: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """``reference: '{case}'`` adopts, per case, the grid of that case's own entry: two cases whose + Ref grids differ land each on their own, not both on a memoized first.""" + monkeypatch.chdir(cohort) + grids = {"P000": ((1.0, 1.2, 0.8), (5.0, -3.0, 2.0)), "P001": ((2.0, 0.7, 1.1), (-8.0, 4.0, 0.5))} + for case, (spacing, origin) in grids.items(): + reference = sitk.GetImageFromArray(np.zeros((5, 6, 7), dtype=np.float32)) + reference.SetSpacing(spacing) + reference.SetOrigin(origin) + sitk.WriteImage(reference, str(cohort / "Raw" / case / "Ref.mha")) + api.transform( + "PER_CASE", + "./Raw:mha", + {"CT": {"Moved": [Resample(reference="{case}", reference_group="Ref"), Write(dataset="./Moved:mha")]}}, + transforms_dir=cohort / "Transforms", + on_fallback="error", + quiet=True, + ) + for case, (spacing, origin) in grids.items(): + moved = sitk.ReadImage(str(cohort / "Moved" / case / "Moved.mha")) + assert moved.GetSpacing() == pytest.approx(spacing) + assert moved.GetOrigin() == pytest.approx(origin) + assert moved.GetSize() == (7, 6, 5) + + +# --------------------------------------------------------------------------- uncertainty vocabulary + + +def test_std_reduction_matches_torch_incrementally() -> None: + rng = np.random.default_rng(3) + members = [torch.from_numpy(rng.normal(size=(1, 4, 5, 6)).astype(np.float32)) for _ in range(5)] + expected = torch.stack(members).std(0) + + torch.testing.assert_close(Std()(list(members)), expected) + + incremental = Std() + incremental.start() + for member in members: + incremental.accumulate(member) + torch.testing.assert_close(incremental.finalize(), expected) + + +def test_std_of_a_single_case_is_zero() -> None: + member = torch.ones(1, 2, 3) + assert Std()([member]).abs().max() == 0.0 + + +def test_magnitude_is_the_channel_norm_and_pointwise() -> None: + from konfai.data.patching import LocalityKind + from konfai.utils.dataset import Attribute + + field = torch.tensor([[[3.0]], [[4.0]]]) + stage = Magnitude() + torch.testing.assert_close(stage("case", field, Attribute()), torch.tensor([[[5.0]]])) + assert stage.patch_locality(Attribute()).kind is LocalityKind.POINTWISE + + +# ------------------------------------------------------------------------------------- config tree + + +def test_a_config_tree_must_hold_the_workflow_root() -> None: + from konfai.utils.runtime import _materialized_config + + with pytest.raises(ConfigError, match="Transformer"): + _materialized_config({"Trainer": {}}, "Transformer") + path = _materialized_config({"Transformer": {"name": "X"}}, "Transformer") + assert path.is_file() From 79ef55730684b89e1c35db8d4176cba11aec6fea Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 00:29:19 +0200 Subject: [PATCH 18/39] feat(impact-reg): every derivation through konfai's own engine The orchestrator stages symlinked cohorts and calls konfai.api: it holds no volume in RAM and resamples nothing by hand. - The moved images are ONE streamed run for the whole cohort: Resample adopts each case's own DVF grid (reference '{case}'), reads the field as the map (field_group), and Write lands Moved beside the DVF. --max-displacement bounds the window a streamed slab reads; 'auto' reads the bound a field recorded (OME-Zarr fields carry one) and falls back to whole-volume with the reason in the plan. - The ensemble average is Reduce(Mean) over members-as-cases, per case: grid strict verifies the shared-grid claim, the fold is incremental, and the written store keeps its displacement-field declaration. - uncertainty is Magnitude -> Reduce(Std) -> Write: no stacked volume, no preset app; the preset argument stays accepted and unused. - evaluate's image and seg warps are the same Resample with the stored transform staged as a group the engine decodes; landmarks still open the transform in-process, a few points. - Transform.h5 stays SimpleITK: the format itself carries the whole field, which is why it is gated behind fields_only. The CLI prints a KonfAIError's message and remedy and exits 1. --- apps/impact_reg/impact_reg_konfai/cli.py | 31 ++ .../impact_reg_konfai/impact_reg.py | 407 ++++++++++-------- .../tests/unit/test_displacement_field_io.py | 82 ++-- .../tests/unit/test_orchestration.py | 11 +- .../tests/unit/test_tmp_dir_forwarding.py | 26 +- 5 files changed, 324 insertions(+), 233 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index bdb425d9..5c2945f9 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -24,11 +24,19 @@ """ import argparse +import sys from pathlib import Path from impact_reg_konfai.impact_reg import ImpactRegKonfAIApp, get_available_presets +def _max_displacement(value: str) -> float | str: + """``auto`` or a distance in world units — the window bound a streamed field read needs.""" + if value.strip().lower() == "auto": + return "auto" + return float(value) + + def _paths(value: str) -> Path: return Path(value).resolve() @@ -134,6 +142,16 @@ def main() -> None: help="Write the displacement fields only: skip the moved image and Transform.h5, both derived " "from the field. For a caller that composes the field itself and would delete them.", ) + reg.add_argument( + "--max-displacement", + "--max_displacement", + dest="max_displacement", + type=_max_displacement, + default="auto", + help="Bound (world units) on the field, sizing what a streamed moved-image slab reads. 'auto' " + "reads the bound the field recorded (OME-Zarr fields carry one) and falls back to whole-volume " + "-- with the reason in the plan -- when none is recorded.", + ) _add_device(reg) _add_tmp_dir(reg) @@ -201,6 +219,18 @@ def main() -> None: download=getattr(args, "download", False), force_update=getattr(args, "force_update", False) ) + # konfai's Python API raises designed refusals (message + remedy); the CLI's job is to print + # them and exit 1 -- the same contract the konfai CLI itself offers. + from konfai.utils.errors import KonfAIError + + try: + _dispatch(args, app, ev) + except KonfAIError as error: + print(str(error).strip(), file=sys.stderr) + sys.exit(1) + + +def _dispatch(args: argparse.Namespace, app: ImpactRegKonfAIApp, ev: argparse.ArgumentParser) -> None: if args.command == "register": gpu = [] if args.cpu is not None else args.gpu app.register( @@ -218,6 +248,7 @@ def main() -> None: config_overrides=args.config_overrides, tmp_dir=args.tmp_dir, fields_only=args.fields_only, + max_displacement=args.max_displacement, ) elif args.command == "eval": diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 9fbd0eb5..dfd09b39 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -111,6 +111,7 @@ def _output_path(dest_dir: Path, stem: str, suffixes: str) -> Path: ``DVF.ome.zarr`` behind, and discovery is by stem: ``_find_output`` takes the first match, which sorts to the stale one. Only the current run's output is left standing. """ + dest_dir.mkdir(parents=True, exist_ok=True) for stale in [p for p in dest_dir.iterdir() if p.name.startswith(f"{stem}.")]: shutil.rmtree(stale) if stale.is_dir() else stale.unlink() return dest_dir / (stem + suffixes) @@ -172,102 +173,53 @@ def _copy_output(src: Path, dest_dir: Path, stem: str) -> Path: return dest -def _write_displacement_field(field: sitk.Image, dest: Path) -> None: - """Write a field in the form ``dest`` names, so a derived field matches the ones it came from.""" - if "".join(dest.suffixes).endswith(".ome.zarr"): - from konfai.utils.dataset import image_to_data - from konfai.utils.ome_zarr import write_ome_zarr - - # image_to_data yields the channel-first array and the Origin/Spacing/Direction the store must - # carry -- the same encoding the Dataset backend writes, so the field round-trips through either. - data, attributes = image_to_data(field) - write_ome_zarr( - dest, - data, - spacing=field.GetSpacing(), - origin=field.GetOrigin(), - attributes=dict(attributes), - displacement_field=True, - ) - else: - sitk.WriteImage(field, str(dest)) - - -def _read_image(path: Path, work: Path) -> sitk.Image: - """A volume, read through ``Dataset`` — after building the dataset it needs. - - A dataset is a root of CASES holding GROUPS. A path on the command line is neither, so it is linked - into that layout first, exactly as konfai-apps stages its own inputs. Pretending instead that the - parent directory is a root and the file a group -- what this used to do -- makes detection probe one - level too deep inside a store, and is simply wrong for a single-store backend like h5, where the - file IS the dataset rather than an entry in one. - - Built properly, no format is named: konfai has a case to probe and detects the backend itself, which - is why this reads an ITK file, an OME-Zarr store or anything else konfai supports without a branch - here enumerating them. - """ - from konfai.utils.dataset import Dataset +def _stage(root: Path, case: str, group: str, source: Path) -> None: + """Link one cohort entry — ``root//`` → ``source``. No bytes move. - root = Path(tempfile.mkdtemp(prefix="entry_", dir=str(work))) - case = root / "P000" - case.mkdir() - (case / f"Entry{''.join(path.suffixes)}").symlink_to(path.resolve()) - return Dataset(root, "").read_image("Entry", "P000") - - -def _write_case_entry(image: sitk.Image, case_out: Path, group: str, file_format: str) -> None: - """Write one group of one case, through ``Dataset``. - - THIS side really is a dataset, and always was: ``//.`` is a root of cases - holding groups, which is why konfai can read it back without being told a format. So the write goes - through the layer that owns that layout instead of composing the path by hand -- and writing, unlike - reading, does need a format named, because there is nothing on disk yet to detect one from. + A dataset is a root of CASES holding GROUPS; paths from the command line (or another run's + output tree) become one by symlink, exactly as konfai-apps stages its own inputs. Entries + resolve by name whatever their form — an ITK file, an OME-Zarr store, a stored transform. """ - from konfai.utils.dataset import Dataset - - Dataset(case_out.parent, file_format).write(group, case_out.name, image) - - -def _derive_moved( - moving_image: Path, dvf_path: Path, dest_dir: Path, work: Path, field: sitk.Image | None = None -) -> Path: - """The moved image, resampled from the moving through the displacement field. - - A preset that emits only a field is complete: the moved image IS that field applied to the moving, - so deriving it belongs to the orchestrator rather than being a second output every preset has to - remember to declare. - - THE FORMAT FOLLOWS THE FIELD, not the moving. Measured: a preset fed ``.mha`` inputs writes its - field as ``.ome.zarr`` when that is what it declares, so the input's form says nothing about the - output's -- the field is the only thing the preset committed to, so the derived moved matches it. - The moving is read through the dataset layer either way; ``sitk.ReadImage``, what this replaces, - cannot open a store at all, which is why the ensemble path failed on OME-Zarr inputs. - - Resampled through SimpleITK for the reason konfai's own ``ResampleTransform`` gives: the stored - displacement is in world (x, y, z) units, and adding it onto a (z, y, x) voxel grid by hand - transposes the axes and reads millimetres as voxels. The output grid is the field's own -- a - displacement field is defined ON the fixed grid, so that is where the moved image belongs. + case_dir = root / case + case_dir.mkdir(parents=True, exist_ok=True) + (case_dir / (group + "".join(source.suffixes))).symlink_to(source.resolve()) + + +def _the_output(dest_dir: Path, stem: str) -> Path: + """The single output named ``stem`` in ``dest_dir`` — one, exactly.""" + matches = sorted(dest_dir.glob(f"{stem}.*")) + if len(matches) != 1: + raise FileNotFoundError(f"Expected exactly one '{stem}' under {dest_dir}, found {len(matches)}.") + return matches[0] + + +def _run_transform( + name: str, + root: Path, + chains: dict, + work: Path, + gpu: list[int], + cpu: int | None, + quiet: bool, +) -> None: + """One streamed TRANSFORM through konfai's Python API, its workspace under ``work``. + + Every derivation runs here: the plan prices each case and routes it (stream, load, + whole-volume — with the reason), so the orchestrator never holds a volume in RAM and never + resamples by hand. A designed refusal raises ``KonfAIError``; the CLI prints it. """ - if field is None: - field = read_displacement_field(dvf_path) - # Read the grid off the field BEFORE the transform takes it: DisplacementFieldTransform assumes - # ownership of the image it is given and leaves it empty behind. - size, origin = field.GetSize(), field.GetOrigin() - spacing, direction = field.GetSpacing(), field.GetDirection() - transform = sitk.DisplacementFieldTransform(field) - moving = _read_image(moving_image, work) - moved = sitk.Resample( - moving, size, transform, sitk.sitkLinear, origin, spacing, direction, 0.0, moving.GetPixelID() + from konfai import api + + api.transform( + name, + f"{root}:mha", # the format token is only the FIRST read candidate; entries resolve by name + chains, + gpu=list(gpu), + cpu=cpu or 1, + quiet=quiet, + overwrite=True, + transforms_dir=work / "Workspaces", ) - suffixes = "".join(dvf_path.suffixes) - dest = _output_path(dest_dir, "Moved", suffixes) - _write_case_entry(moved, dest_dir, "Moved", _FORMATS.get(suffixes.lower(), suffixes.lstrip("."))) - return dest - - -def _displacement_transform(dvf_path: Path) -> sitk.Transform: - """Read a displacement field (3-component, fixed grid) as a SimpleITK transform.""" - return sitk.DisplacementFieldTransform(read_displacement_field(dvf_path)) def _neutral_mask(out_path: Path) -> Path: @@ -376,6 +328,7 @@ def register( config_overrides: list[str] | None = None, tmp_dir: Path | None = None, fields_only: bool = False, + max_displacement: float | str = "auto", ) -> None: """Register every case with the selected presets and ensemble their DVFs. @@ -393,6 +346,11 @@ def register( full-size rewrite -- worth it for a caller that wants a registration to look at, waste for one that composes the field with another and derives its own. A caller that reads only the fields should be able to say so rather than pay for outputs it deletes. + + ``max_displacement`` bounds the window the field is read from when the moved image streams: + ``auto`` reads the bound a field recorded (OME-Zarr fields carry one) and falls back to the + whole volume -- with the reason in the plan -- when none is recorded; a distance in world + units declares it outright and is checked against every region actually read. """ # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs @@ -433,7 +391,7 @@ def register( "the moved image is derived per case and needs the two to line up." ) - for case, moving_image in zip(cases, moving_units, strict=True): + for case in cases: # konfai-apps already named the cases; reusing its names keeps the two layers' notion of # a case identical instead of renumbering from the command line and hoping they agree. case_out = output / case @@ -451,44 +409,118 @@ def register( dvf_paths.append(dvf) if len(presets) == 1: - dvf_out = _copy_output(dvf_paths[0], case_out, "DVF") - if not fields_only: - _derive_moved(moving_image, dvf_out, case_out, work) + _copy_output(dvf_paths[0], case_out, "DVF") else: - # Ensemble: average the presets' displacement fields (all on the fixed grid) and warp the - # moving image once with that averaged field — the one output no single preset produced. - avg_dvf = self._average_displacement(dvf_paths) - dvf_out = _output_path(case_out, "DVF", "".join(dvf_paths[0].suffixes)) - _write_displacement_field(avg_dvf, dvf_out) - if not fields_only: - # Through the same derivation as the single-preset path, which reads the moving - # in either form: the sitk.ReadImage this replaces cannot open a store at all, - # so an ensemble of OME-Zarr inputs failed here and nowhere else. - _derive_moved( - moving_image, dvf_out, case_out, work, field=sitk.Cast(avg_dvf, sitk.sitkVectorFloat64) - ) - - if not fields_only: - # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid - # displacement field as a SimpleITK transform. Another full-size write of the same - # voxels, which is why it goes with the moved image rather than being unconditional. - sitk.WriteTransform(_displacement_transform(dvf_out), str(case_out / "Transform.h5")) + # Ensemble: fold the presets' fields (all on the fixed grid -- and Reduce VERIFIES + # the claim) into the averaged DVF, the one output no single preset produced. + # Streamed: the fold is incremental, so the peak is one accumulator plus the + # member being read, whatever the size of the ensemble. + self._ensemble_mean(case, presets, dvf_paths, output, work, gpu, cpu, quiet) + + if not fields_only: + # Every moved image in ONE streamed run: Resample adopts, per case, the grid of that + # case's own DVF (a field is defined ON the fixed grid) and reads the field as the + # map -- one interpolation, slab by slab, the whole cohort under one plan. + self._derive_moved( + dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet, max_displacement + ) + for case in cases: + # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid field + # as a SimpleITK transform. Inherently whole -- the .h5 format carries the full + # field -- which is why it goes with the moved image rather than being + # unconditional. + transform = sitk.DisplacementFieldTransform( + read_displacement_field(_the_output(output / case, "DVF")) + ) + sitk.WriteTransform(transform, str(output / case / "Transform.h5")) finally: shutil.rmtree(work, ignore_errors=True) - def _average_displacement(self, dvf_paths: list[Path]) -> sitk.Image: - """Average several presets' displacement fields (all on the same fixed grid) into one field. + def _ensemble_mean( + self, + case: str, + presets: list[str], + dvf_paths: list[Path], + output: Path, + work: Path, + gpu: list[int], + cpu: int | None, + quiet: bool, + ) -> None: + """Average one case's preset fields into ``//DVF`` — Reduce(Mean), streamed.""" + from konfai.data.transform import Reduce, Write + + root = work / f"ensemble_{case}" + for preset, dvf in zip(presets, dvf_paths, strict=True): + _stage(root, preset, "DVF", dvf) + suffixes = "".join(dvf_paths[0].suffixes) + _output_path(output / case, "DVF", suffixes) # drop a stale other-form DVF before writing + _run_transform( + f"impact_reg_ensemble_{case}", + root, + { + "DVF": { + "DVF": [ + Reduce(operator="Mean", output=case, grid="strict"), + Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), + ] + } + }, + work, + gpu, + cpu, + quiet, + ) + + def _derive_moved( + self, + cases: dict[str, Path], + output: Path, + work: Path, + gpu: list[int], + cpu: int | None, + quiet: bool, + max_displacement: float | str, + ) -> None: + """The moved images, resampled from each moving through ITS displacement field — one run. - A running sum keeps memory flat in the number of members: a few field-sized buffers are live - at any instant, whatever the size of the ensemble. + A preset that emits only a field is complete: the moved image IS that field applied to the + moving, so deriving it belongs to this layer. THE GRID AND THE FORMAT FOLLOW THE FIELD, not + the moving: a displacement field is defined ON the fixed grid, so ``reference: '{case}'`` + adopts each case's own DVF grid, and the field is the map (``field_group``) — one + interpolation, streamed when the field's bound allows, whole-volume with the reason when not. """ - reference = read_displacement_field(dvf_paths[0]) - total = sitk.GetArrayFromImage(reference) - for path in dvf_paths[1:]: - total += sitk.GetArrayFromImage(read_displacement_field(path)) - avg = sitk.GetImageFromArray(total / len(dvf_paths), isVector=True) - avg.CopyInformation(reference) - return avg + from konfai.data.transform import Resample, Write + + root = work / "moved_stage" + suffixes = "" + for case, moving in cases.items(): + dvf = _the_output(output / case, "DVF") + _stage(root, case, "Moving", moving) + _stage(root, case, "DVF", dvf) + suffixes = "".join(dvf.suffixes) + _output_path(output / case, "Moved", suffixes) # drop a stale other-form Moved + _run_transform( + "impact_reg_moved", + root, + { + "Moving": { + "Moved": [ + Resample( + reference="{case}", + reference_group="DVF", + field_group="DVF", + max_displacement=max_displacement, + ), + Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), + ] + } + }, + work, + gpu, + cpu, + quiet, + ) # ------------------------------------------------------------------ evaluate @@ -527,16 +559,13 @@ def evaluate( n_cases = max(len(fixed_images), len(gt_fixed_seg), len(gt_fixed_fid)) for index in range(n_cases): transform_path = transforms[index] if index < len(transforms) else None - transform = sitk.ReadTransform(str(transform_path)) if transform_path else sitk.Transform() eval_out = output / f"P{index:03d}" / "Evaluation" work = _work_dir(tmp_dir, "impact_reg_eval_") try: # Image: moving resampled onto the fixed grid vs fixed (MAE). Mask is optional. if index < len(fixed_images) and index < len(moving_images): - fixed = _read_image(fixed_images[index], work) - moved = work / "moved_image.nii.gz" - sitk.WriteImage( - sitk.Resample(_read_image(moving_images[index], work), fixed, transform), str(moved) + moved = self._warp_onto_fixed( + work, "img", fixed_images[index], moving_images[index], transform_path, gpu, cpu, quiet ) app.evaluate( inputs=[[fixed_images[index]]], @@ -550,15 +579,10 @@ def evaluate( tmp_dir=work, ) - # Segmentation: moving seg warped onto fixed vs fixed seg (Dice). + # Segmentation: moving seg warped onto fixed vs fixed seg (Dice), nearest-neighbour. if index < len(gt_fixed_seg) and index < len(gt_moving_seg): - fixed_seg = _read_image(gt_fixed_seg[index], work) - moved_seg = work / "moved_seg.nii.gz" - sitk.WriteImage( - sitk.Resample( - _read_image(gt_moving_seg[index], work), fixed_seg, transform, sitk.sitkNearestNeighbor - ), - str(moved_seg), + moved_seg = self._warp_onto_fixed( + work, "seg", gt_fixed_seg[index], gt_moving_seg[index], transform_path, gpu, cpu, quiet ) app.evaluate( inputs=[[gt_fixed_seg[index]]], @@ -574,10 +598,12 @@ def evaluate( # Landmarks (TRE): the transform is defined on the fixed grid and maps fixed->moving, so the # fixed fiducials are displaced by it into moving space and compared against the moving fiducials # there (the standard warped-keypoints convention; no field inversion needed). With no transform - # the raw fiducials are compared, measuring the initial misalignment. + # the raw fiducials are compared, measuring the initial misalignment. Landmarks are a few + # points, so this is the one place the transform is opened in this process. if index < len(gt_fixed_fid) and index < len(gt_moving_fid): fixed_points = read_landmarks(gt_fixed_fid[index]) if transform_path is not None: + transform = sitk.ReadTransform(str(transform_path)) fixed_points = apply_to_data_transform(fixed_points, {transform: False}) moved_fid = work / "moved_fid.fcsv" write_landmarks(fixed_points, moved_fid) @@ -594,6 +620,47 @@ def evaluate( finally: shutil.rmtree(work, ignore_errors=True) + def _warp_onto_fixed( + self, + work: Path, + kind: str, + fixed: Path, + moving: Path, + transform_path: Path | None, + gpu: list[int], + cpu: int | None, + quiet: bool, + ) -> Path: + """The moving warped onto the fixed grid — one streamed Resample; nearest for a ``seg``. + + ``reference: '{case}'`` adopts the fixed grid; the transform, when given, is staged as a + stored-transform group konfai decodes itself (rigid, affine, spline, field or composite) — + with its exact affine box or its coefficient-derived bound sizing what each slab reads. + With no transform the map is the identity and this is a change of grid alone. + """ + from konfai.data.transform import Resample, Write + + root = work / f"warp_{kind}" + _stage(root, "P000", "Fixed", fixed) + _stage(root, "P000", "Moving", moving) + resample: dict[str, object] = {"reference": "{case}", "reference_group": "Fixed"} + if kind == "seg": + resample["interpolation"] = "nearest" + if transform_path is not None: + _stage(root, "P000", "Reg", transform_path) + resample["transforms"] = {"Reg": False} + out_root = work / f"moved_{kind}" + _run_transform( + f"impact_reg_eval_{kind}", + root, + {"Moving": {"Moved": [Resample(**resample), Write(dataset=f"{out_root}:mha")]}}, + work, + gpu, + cpu, + quiet, + ) + return out_root / "P000" / "Moved.mha" + # --------------------------------------------------------------- uncertainty def uncertainty( @@ -608,45 +675,39 @@ def uncertainty( ) -> None: """Estimate registration uncertainty as the voxel-wise spread of an ensemble of displacement fields. - The per-preset displacement fields are stacked into one multi-component volume (samples as - components, vector components as the leading image axis) and handed to the preset's generic - ``Uncertainty.yml`` workflow (``konfai-apps uncertainty``: ``Norm`` magnitude then - ``StandardDeviation`` over the ensemble). + Each member's magnitude, then the standard deviation across members: ``Magnitude`` is + pointwise and ``Reduce(Std)`` folds with running moments, so no member is ever whole in RAM + whatever the size of the ensemble. ``preset`` is accepted for CLI compatibility and unused — + measuring a spread needs no preset app. Writes ``/uncertainty/Uncertainty.``. """ + del preset if len(dvfs) < 2: raise ValueError("Uncertainty needs at least two ensemble displacement fields.") work = _work_dir(tmp_dir, "impact_reg_unc_") try: - reference = read_displacement_field(dvfs[0]) - rank = reference.GetDimension() - stack = sitk.GetImageFromArray( - np.stack([sitk.GetArrayFromImage(read_displacement_field(p)) for p in dvfs], axis=-1), - isVector=True, + from konfai.data.transform import Magnitude, Reduce, Write + + root = work / "members" + members = _units(list(dvfs)) + for index, dvf in enumerate(members): + _stage(root, f"M{index:03d}", "DVF", dvf) + suffixes = "".join(members[0].suffixes) + _run_transform( + "impact_reg_uncertainty", + root, + { + "DVF": { + "Uncertainty": [ + Magnitude(), + Reduce(operator="Std", output="uncertainty", grid="strict"), + Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), + ] + } + }, + work, + gpu, + cpu, + quiet, ) - # The extra leading image axis holds the vector components (dropped by ``Norm``); the real - # fixed-grid geometry lives on the remaining axes so the uncertainty map stays aligned. - stack.SetOrigin((0.0, *reference.GetOrigin())) - stack.SetSpacing((1.0, *reference.GetSpacing())) - direction = np.eye(rank + 1) - direction[1:, 1:] = np.asarray(reference.GetDirection()).reshape(rank, rank) - stack.SetDirection(direction.flatten()) - sitk.WriteImage(stack, str(work / "DVFs.mha")) - - # Same workspace hand-off as _infer_preset: without it konfai-apps auto-creates one under - # TMPDIR and stages Uncertainties there before copying it into -o, which is the staging - # this option exists to place. `work` is ours and already sits wherever tmp_dir asked for. - command = ["konfai-apps", "uncertainty", _app_id(preset), "-i", str(work / "DVFs.mha"), - "-o", str(output), "--tmp-dir", str(work)] - if gpu: - command += ["--gpu", *(str(g) for g in gpu)] - elif cpu is not None: - command += ["--cpu", str(cpu)] - if quiet: - command.append("--quiet") - if self._download: - command.append("--download") - if self._force_update: - command.append("--force_update") - subprocess.run(command, check=True) # nosec B603 finally: shutil.rmtree(work, ignore_errors=True) diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index d2854b0c..d8efdb30 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -31,17 +31,34 @@ pytest.importorskip("ngff_zarr") from impact_reg_konfai.impact_reg import ( # noqa: E402 + ImpactRegKonfAIApp, _copy_output, - _displacement_transform, _find_outputs, _output_path, - _read_image, - _write_displacement_field, ) from konfai.utils.errors import TransformError # noqa: E402 from konfai.utils.ITK import read_displacement_field # noqa: E402 from konfai.utils.ome_zarr import _zarr_v3_available, is_displacement_field, write_ome_zarr # noqa: E402 + +def _write_displacement_field(field: "sitk.Image", dest: Path) -> None: + """Fixture-builder: a field in either form, as a preset app would have produced it.""" + if "".join(dest.suffixes).endswith(".ome.zarr"): + from konfai.utils.dataset import image_to_data + + data, attributes = image_to_data(field) + write_ome_zarr( + dest, + data, + spacing=field.GetSpacing(), + origin=field.GetOrigin(), + attributes=dict(attributes), + displacement_field=True, + ) + else: + sitk.WriteImage(field, str(dest)) + + # The OME-Zarr side of these tests writes an RFC-5 field, a zarr v3 store that zarr 2.x # (Python 3.10) cannot write. pytestmark = pytest.mark.skipif( @@ -203,10 +220,26 @@ def test_rerunning_in_the_other_form_leaves_one_output(tmp_path: Path) -> None: assert _find_outputs(destination, "DVF")[destination.name].name == "DVF.ome.zarr" -def test_store_written_by_the_orchestrator_is_a_declared_field(tmp_path: Path) -> None: - """An averaged ensemble field is written in the members' form, and stays a declared field.""" - store = _write_store(tmp_path / "DVF.ome.zarr", _field()) - assert is_displacement_field(store) +def test_ensemble_field_written_by_the_orchestrator_is_a_declared_field(tmp_path: Path) -> None: + """The averaged DVF is folded by Reduce(Mean) and written through konfai's Write, in the + members' form — and must stay a DECLARED field: ``read_displacement_field`` refuses an + undeclared 3-channel store, so dropping the declaration would break Transform.h5 and + ``evaluate`` right after a successful register. The values are the voxel-wise mean, on the + members' geometry.""" + members = [] + for index in (1, 2): + (tmp_path / f"m{index}").mkdir() + members.append(_write_store(tmp_path / f"m{index}" / "DVF.ome.zarr", _field(index))) + output, work = tmp_path / "out", tmp_path / "work" + work.mkdir() + + ImpactRegKonfAIApp()._ensemble_mean("P000", ["a", "b"], members, output, work, [], 1, True) + + out = output / "P000" / "DVF.ome.zarr" + assert is_displacement_field(out) + averaged = sitk.GetArrayFromImage(read_displacement_field(out)) + expected = (sitk.GetArrayFromImage(_field(1)) + sitk.GetArrayFromImage(_field(2))) / 2 + np.testing.assert_allclose(averaged, expected, rtol=1e-5, atol=1e-5) @pytest.mark.parametrize("suffix", [".mha", ".ome.zarr"]) @@ -219,42 +252,9 @@ def test_transform_reads_back_identically_from_either_form(tmp_path: Path, suffi original = _field() _write_displacement_field(original, tmp_path / f"DVF{suffix}") - restored = _displacement_transform(tmp_path / f"DVF{suffix}") + restored = sitk.DisplacementFieldTransform(read_displacement_field(tmp_path / f"DVF{suffix}")) assert isinstance(restored, sitk.DisplacementFieldTransform) reference = sitk.DisplacementFieldTransform(sitk.Image(original)) for point in ((9.0, -1.0, 12.0), (7.5, -2.5, 11.0)): assert restored.TransformPoint(point) == pytest.approx(reference.TransformPoint(point)) - - -@pytest.mark.skipif(not _zarr_v3_available(), reason="writing an OME-Zarr store needs zarr 3") -def test_read_image_opens_an_ome_zarr_store(tmp_path) -> None: - """An ordinary image reads back from a store, not only from an ITK file. - - Read through ``Dataset`` with no format named: the path is staged into a real case/group layout - first, so konfai detects the backend itself. ``sitk.ReadImage`` cannot open a directory, which is - what made the ensemble path unusable with OME-Zarr inputs while the single-preset path — which - reuses the model's output and never re-reads — worked fine. - """ - volume = np.arange(4 * 5 * 6, dtype=np.float32).reshape(4, 5, 6) - store = tmp_path / "moving.ome.zarr" - write_ome_zarr(store, volume[np.newaxis], spacing=(1.5, 2.0, 2.5), origin=(3.0, -1.0, 0.5)) - - work = tmp_path / "work" - work.mkdir() - image = _read_image(store, work) - - assert image.GetSpacing() == pytest.approx((1.5, 2.0, 2.5)) - assert image.GetOrigin() == pytest.approx((3.0, -1.0, 0.5)) - np.testing.assert_allclose(sitk.GetArrayFromImage(image), volume) - - -def test_read_image_still_opens_an_itk_file(tmp_path) -> None: - """The same call, on a plain ITK file: one reader, no branch on the form.""" - volume = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) - path = tmp_path / "moving.mha" - sitk.WriteImage(sitk.GetImageFromArray(volume), str(path)) - - work = tmp_path / "work" - work.mkdir() - np.testing.assert_allclose(sitk.GetArrayFromImage(_read_image(path, work)), volume) diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index 8ad4d7ce..94bb24ed 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -79,7 +79,9 @@ def _write_dvf(path: Path, vector, reference: sitk.Image) -> Path: return path -def test_average_displacement_is_the_voxelwise_mean_with_reference_geometry(tmp_path: Path) -> None: +def test_ensemble_mean_is_the_voxelwise_mean_with_reference_geometry(tmp_path: Path) -> None: + """Reduce(Mean) over members-as-cases: the averaged DVF lands at //DVF, on the + members' shared grid — which ``grid: strict`` verified rather than assumed.""" reference = sitk.GetImageFromArray(np.zeros((6, 6, 6), dtype=np.float32)) reference.SetSpacing((1.5, 1.5, 1.5)) reference.SetOrigin((3.0, -2.0, 1.0)) @@ -87,7 +89,12 @@ def test_average_displacement_is_the_voxelwise_mean_with_reference_geometry(tmp_ _write_dvf(tmp_path / "a.mha", (1.0, 0.0, 0.0), reference), _write_dvf(tmp_path / "b.mha", (3.0, 2.0, -4.0), reference), ] - avg = reg.ImpactRegKonfAIApp()._average_displacement(paths) + output, work = tmp_path / "out", tmp_path / "work" + work.mkdir() + + reg.ImpactRegKonfAIApp()._ensemble_mean("P000", ["a", "b"], paths, output, work, [], 1, True) + + avg = sitk.ReadImage(str(output / "P000" / "DVF.mha")) field = sitk.GetArrayFromImage(avg) np.testing.assert_allclose(field[0, 0, 0], (2.0, 1.0, -2.0), atol=1e-6) assert avg.GetSpacing() == pytest.approx((1.5, 1.5, 1.5)) diff --git a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py index 128636d5..7b244b6a 100644 --- a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py +++ b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py @@ -57,22 +57,16 @@ def fake_run(command, **kwargs): work = tmp_path / "work" work.mkdir() app = ImpactRegKonfAIApp() - app._infer_preset( - "FireANTs_SyN", [tmp_path / "f.mha"], [tmp_path / "m.mha"], [], [], 1, work, [], None, True - ) + app._infer_preset("FireANTs_SyN", [tmp_path / "f.mha"], [tmp_path / "m.mha"], [], [], 1, work, [], None, True) assert len(captured) == 1 assert _tmp_dir_value(captured[0]) == str(work / "FireANTs_SyN") -def test_uncertainty_forwards_the_workspace(tmp_path: Path, monkeypatch, write_preset_output) -> None: - """``konfai-apps uncertainty`` stages inside the caller's tmp_dir, not under the system TMPDIR.""" - captured: list[list[str]] = [] - monkeypatch.setattr( - "impact_reg_konfai.impact_reg.subprocess.run", - lambda command, **kwargs: captured.append(list(command)), - ) - +def test_uncertainty_stages_inside_the_callers_tmp_dir(tmp_path: Path, write_preset_output) -> None: + """The staging and the run workspaces live in a private directory INSIDE the caller's tmp_dir -- + never under the system TMPDIR -- and the caller's directory is left standing, emptied, when the + run is done. The spread map is the one deliverable, under /uncertainty/.""" _, first = write_preset_output(tmp_path / "a") _, second = write_preset_output(tmp_path / "b") staging = tmp_path / "staging" @@ -85,9 +79,7 @@ def test_uncertainty_forwards_the_workspace(tmp_path: Path, monkeypatch, write_p tmp_dir=staging, ) - assert len(captured) == 1 - # The workspace is the private directory _work_dir made INSIDE the caller's tmp_dir -- never the - # caller's own directory, which the command must leave standing. - forwarded = Path(_tmp_dir_value(captured[0])) - assert forwarded.parent == staging - assert forwarded != staging + assert staging.is_dir() + assert list(staging.iterdir()) == [] + spread = sorted((tmp_path / "out" / "uncertainty").iterdir()) + assert [path.name for path in spread] == ["Uncertainty.mha"] From a7b429f5d257d8e809d01bd7e1af3b5284896366 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 01:09:06 +0200 Subject: [PATCH 19/39] docs(mcp): plan_transform's verdict list gains LOAD The TRANSFORM plan answers LOAD -- the case fits the budget and streaming would reread the source past its worth -- and the tool description enumerated every verdict but that one. The generated tool reference follows. --- .../skills/konfai-experiments/references/tool-reference.md | 2 +- konfai-mcp/konfai_mcp/guide.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.claude/skills/konfai-experiments/references/tool-reference.md b/.claude/skills/konfai-experiments/references/tool-reference.md index 63a53009..975e5e98 100644 --- a/.claude/skills/konfai-experiments/references/tool-reference.md +++ b/.claude/skills/konfai-experiments/references/tool-reference.md @@ -132,7 +132,7 @@ Use to PACKAGE a model trained in the current session (the train-from-scratch br ### `plan_transform` -Use BEFORE run_transform, always: it is the dry run, and it writes no data. This plans every (case, chain) from the session Transform.yml. The plan is a measurement, not an estimate -- it opens and removes a real region-write on each destination -- so its verdict is the one the run will act on: STREAM (bounded memory), WHOLE-VOLUME (the case is assembled whole, with the stage that refused it named), SKIP (already written), REDUCE, REFUSED. Outputs: {ok, report, verdict_counts, budget, needs_attention[], over_budget[]}. A non-empty over_budget means run_transform would refuse before writing anything. Next: fix what needs_attention names, or run_transform. +Use BEFORE run_transform, always: it is the dry run, and it produces none of the deliverable. This plans every (case, chain) from the session Transform.yml. The plan is a measurement, not an estimate -- it opens and removes a real region-write on each destination -- so its verdict is the one the run will act on: STREAM (bounded memory), LOAD (the case fits the budget and streaming would reread the source past its worth -- a cost choice, not a fallback), WHOLE-VOLUME (the case is assembled whole, with the stage that refused it named), SKIP (already written), REDUCE, REFUSED. That probe TOUCHES the output locations, which are the user's own stores: an entry is created and removed, and a single-file store (h5) is created if it did not exist. Outputs: {ok, config_path, world_size, report, verdict_counts, budget, budget_bytes, needs_attention[], over_budget[]}. A non-empty over_budget means run_transform would refuse before writing anything; needs_attention lists at most 50 entries, so trust verdict_counts for the totals. Next: fix what needs_attention names, or run_transform. ### `prepare_dataset_aliases` diff --git a/konfai-mcp/konfai_mcp/guide.py b/konfai-mcp/konfai_mcp/guide.py index 81e787b0..7e180072 100644 --- a/konfai-mcp/konfai_mcp/guide.py +++ b/konfai-mcp/konfai_mcp/guide.py @@ -507,8 +507,9 @@ "Use BEFORE run_transform, always: it is the dry run, and it produces none of the deliverable. " "This plans every (case, chain) from the session Transform.yml. The plan is a measurement, not an " "estimate -- it opens and removes a real region-write on each destination -- so its verdict is the one " - "the run will act on: STREAM (bounded memory), WHOLE-VOLUME (the case is assembled whole, with the stage " - "that refused it named), SKIP (already written), REDUCE, REFUSED. That probe TOUCHES the output " + "the run will act on: STREAM (bounded memory), LOAD (the case fits the budget and streaming would reread " + "the source past its worth -- a cost choice, not a fallback), WHOLE-VOLUME (the case is assembled whole, " + "with the stage that refused it named), SKIP (already written), REDUCE, REFUSED. That probe TOUCHES the output " "locations, which are the user's own stores: an entry is created and removed, and a single-file store " "(h5) is created if it did not exist. " "Outputs: {ok, config_path, world_size, report, verdict_counts, budget, budget_bytes, needs_attention[], " From 59fa47716a1fac6b4d6cd4bb9d0d135e10b1c8bc Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 01:42:57 +0200 Subject: [PATCH 20/39] feat(data): size a field's windows from the field itself at run A field with no declared or recorded bound was a whole-volume answer: the plan could not size what a region must read without reading the field. It streams now. The field window a region samples is its own box, read for sampling regardless -- and the sup of those very values bounds every interpolated displacement in the region (a convex combination cannot exceed the lattice values it blends). The window is memoized so sizing and sampling share one read, and the halo is per region and per component: a quiet slab pays a quiet halo. The plan stays headers-only: the run walks a measured pull (_ReadStagePlan.run_pull) while the estimator prices the declared one -- as if the field were zero when nothing bounds it, and the plan prints that note. A declared or recorded bound keeps the declared windows, whose streamed result is pinned bit-identical to the whole path, and stays checked against every region read; measuring is the route for the field that could not stream at all before. An unreadable entry in the field group keeps its whole-volume answer -- it fails both routes at run. --- apps/impact_reg/impact_reg_konfai/cli.py | 6 +- .../impact_reg_konfai/impact_reg.py | 8 +- docs/source/config_guide/transform.md | 14 +- .../source/reference/components/transforms.md | 2 +- konfai/data/patching.py | 18 ++- konfai/data/transform.py | 131 ++++++++++++++---- tests/unit/test_resample_to_reference.py | 15 +- .../unit/test_transform_locality_contract.py | 16 ++- tests/unit/test_warp.py | 60 +++++++- 9 files changed, 210 insertions(+), 60 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index 5c2945f9..25b0778b 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -148,9 +148,9 @@ def main() -> None: dest="max_displacement", type=_max_displacement, default="auto", - help="Bound (world units) on the field, sizing what a streamed moved-image slab reads. 'auto' " - "reads the bound the field recorded (OME-Zarr fields carry one) and falls back to whole-volume " - "-- with the reason in the plan -- when none is recorded.", + help="Optional bound (world units) on the field: the moved image streams either way, each slab " + "sized from the field values it reads. A bound ('auto' reads the one OME-Zarr fields record) " + "lets the plan price the reads exactly, and a declared one is checked against every region read.", ) _add_device(reg) _add_tmp_dir(reg) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index dfd09b39..6ebf404d 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -347,10 +347,10 @@ def register( that composes the field with another and derives its own. A caller that reads only the fields should be able to say so rather than pay for outputs it deletes. - ``max_displacement`` bounds the window the field is read from when the moved image streams: - ``auto`` reads the bound a field recorded (OME-Zarr fields carry one) and falls back to the - whole volume -- with the reason in the plan -- when none is recorded; a distance in world - units declares it outright and is checked against every region actually read. + ``max_displacement`` is optional: the moved image streams either way, each slab's source + window sized from the field values read for sampling. ``auto`` reads the bound a field + recorded (OME-Zarr fields carry one) so the plan prices the reads exactly; a distance in + world units declares that bound outright and is checked against every region actually read. """ # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index b56c3876..758daaa1 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -373,12 +373,16 @@ 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 +`max_displacement` is **optional**. The field window a region samples is its +own box, read for sampling regardless — and the sup of the values just read +bounds that region's source pull, so each slab pays exactly the halo *its* +displacements require, measured at run from a read the sampler needed anyway. +A declared bound (or `auto`, reading the one KonfAI records on a field it +writes) does two things: it lets the plan **price** the reads exactly — with +no bound the estimate assumes a zero field, and the plan says so — and it 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. +dark rim around the moved anatomy and nothing else. Naming no target grid is the shape update of an atlas build — the field applied on the case's *own* grid — and is the same stage with `reference` left out: diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 6a9fe829..7975c39e 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -97,7 +97,7 @@ until it declares otherwise. | --- | --- | --- | --- | --- | --- | | `Padding` | `F.pad`; updates Origin. `mode` supports `"constant:"`. | `padding=[0,0,0,0,0,0], mode="constant", inverse=True` | **yes** | **yes** | no‡ | | `Crop` | Crop to foreground bounding box; caches the box; updates Origin. | `inverse=True` | **yes** | **yes** (pads back) | **yes** — once the `box` is on the case; the region is the patch translated | -| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk bounds by `max_displacement`, **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, `invert: true` names a spline or a field, or a field has no bound | +| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region — and a declared/recorded `max_displacement` prices the plan exactly and is **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | | `ResampleToResolution` | Deprecated spelling of `Resample: {spacing: ...}`. | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** | | `ResampleToShape` | Deprecated spelling of `Resample: {shape: ...}`. | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** | | `ResampleToReference` | Deprecated spelling of `Resample: {reference: ...}`. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** | **yes** | diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 6ad8b66e..73db817f 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -283,12 +283,17 @@ def __call__(self, target: tuple[slice, ...]) -> list[slice]: class _ReadStagePlan: """One chain stage as the composed streamed read runs it: its declared kind, the spatial shapes on either side, and — for a region stage — the pull map from a region of its output to the region - of its input it is computed from, bound to the case state the stages before it left.""" + of its input it is computed from, bound to the case state the stages before it left. + + ``run_pull``, when set, is the pull the RUN walks instead: a stage that sizes its windows from + the data it reads (a declared field) measures there, while ``pull`` stays headers-only for the + plan's pricing — the estimator must never read a voxel.""" kind: LocalityKind in_shape: tuple[int, ...] out_shape: tuple[int, ...] pull: Callable[[tuple[slice, ...]], list[slice]] | None + run_pull: Callable[[tuple[slice, ...]], list[slice]] | None = None def save_destination(save: Save, default_dataset: Dataset, default_group: str) -> tuple[Dataset, str]: @@ -1828,9 +1833,15 @@ def _plan_read_stage( ) # ORIENTATION / CROP / REGRID: the stage's own remap, on the state the stages before it left. pull = _RemapPull(stage.stream_region_source, list(shape), Attribute(evolved), self.name) + measured = getattr(stage, "measured_region_source", None) + run_pull = ( + _RemapPull(measured, list(shape), Attribute(evolved), self.name) + if measured is not None and getattr(stage, "measures_at_run", False) + else None + ) out = self._stage_out_shape(stage, shape, Attribute(evolved)) stage.write_stream_cache_attribute(evolved, list(shape), self.name) - return _ReadStagePlan(loc.kind, tuple(shape), tuple(out), pull) + return _ReadStagePlan(loc.kind, tuple(shape), tuple(out), pull, run_pull) def _stage_out_shape(self, stage: Stage, shape: list[int], attribute: Attribute) -> list[int]: """The spatial shape one stage folds ``shape`` to — a transform's map or a draw's own. @@ -2924,7 +2935,8 @@ def _replay_streamed_region( spans: list[list[slice]] = [list(target_slices)] for plan in reversed(plans): - spans.append(plan.pull(tuple(spans[-1])) if plan.pull is not None else list(spans[-1])) + pull = plan.run_pull or plan.pull + spans.append(pull(tuple(spans[-1])) if pull is not None else list(spans[-1])) spans.reverse() data_slices = tuple([slice(None)] * n_prefix + spans[0]) diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 2693ca99..3186dd01 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -1197,12 +1197,16 @@ class Resample(TransformInverse): stages resamples twice, and a volume interpolated twice has lost detail the second pass invented no more of -- which is the whole reason an atlas's appearance is rebuilt from native volumes. - IT STREAMS, and what a region reads is known before a voxel is touched. A rigid or affine map is - an exact affine, so the source box of a target region is that region's box mapped through it. A - BSpline and a dense field are values on a grid read through a non-negative kernel that sums to - one, so the sup-norm of those values bounds the displacement at EVERY point -- a theorem, not a - sample of the boundary. A field on disk is bounded by ``max_displacement`` instead, which is then - CHECKED against every region actually read. + IT STREAMS, and what a region reads is known before a voxel of the SOURCE is touched. A rigid + or affine map is an exact affine, so the source box of a target region is that region's box + mapped through it. A BSpline and a dense field are values on a grid read through a non-negative + kernel that sums to one, so the sup-norm of those values bounds the displacement at EVERY point + -- a theorem, not a sample of the boundary. A field on disk is read region by region, and the + window a region samples is its own box: the sup of the values just read bounds that region's + pull, so each slab pays exactly the halo ITS displacements require -- measured at run, from a + read the sampler needs regardless. ``max_displacement`` is optional: declared (or recorded by + the store at write time) it prices the plan exactly and is CHECKED against every region read; + absent, the plan prices the reads as if the field were zero and says so. ``align`` decides where a ``spacing`` or a ``shape`` grid SITS, and it is the one silent choice in the family -- a quarter of a voxel of anatomy, made differently by every library that offers @@ -1218,7 +1222,6 @@ class Resample(TransformInverse): - ``invert: true`` on anything but a rigid or affine map: inverting a spline or a field is a dense solve over the whole grid, and a field solved per region is not the restriction of the field solved once. Store the inverse, or invert it where it is written; - - a field with no ``max_displacement`` to size its region from; - a case that does not meet the target grid anywhere -- judged THROUGH the declared map, so a stored rigid bridging two scanner frames is not mistaken for disjointness. The output would be ``fill`` from edge to edge, and an all-background member is a plausible, wrong @@ -1285,6 +1288,9 @@ def __init__( #: Per case: the geometry keys its header did not carry (see :meth:`Grid.from_header`). self._assumed: dict[str, frozenset[str]] = {} self._stored: dict[str, SpatialStages] = {} + #: 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 self._refusal: str | None = None self._probed = False @@ -1440,12 +1446,17 @@ def _stored_stages(self, name: str) -> SpatialStages: return self._stored[name] def _field_stage(self, name: str, region: Grid) -> DisplacementStage: - """The declared field over ``region``, read on its own grid and no wider. + """The declared field over ``region``, read on its own grid and no wider — once. The field is evaluated at the TARGET's world points, so the window it needs is that region's own world box -- no halo, whatever the displacement is. What the halo sizes is the SOURCE - read, which is a different question answered by the bound. + read, which is a different question, answered from these very values + (:meth:`measured_region_source`) -- memoized here so sizing and sampling share one read. """ + key = (tuple(int(extent) for extent in region.size_zyx), tuple(float(v) for v in np.ravel(region.origin_xyz))) + cached = self._field_window + if cached is not None and cached[0] == name and cached[1] == key: + return cached[2] source = cast("_DisplacementSource", self.displacement) shape, attribute = source.infos(name) spatial = [int(extent) for extent in shape[1:]] @@ -1453,7 +1464,13 @@ def _field_stage(self, name: str, region: Grid) -> DisplacementStage: window = grid.index_window(region.world_box(), margin=1) values = source.read(name, window, len(spatial)) source.check_bound(values, name) - return DisplacementStage(grid.sub_grid(window), values.numpy(), order=1) + 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: + if self._field_window is not None and self._field_window[0] == name: + self._field_window = None def _stages(self, name: str, region: Grid) -> SpatialStages: """The whole map over one target region, in application order.""" @@ -1477,6 +1494,22 @@ def _bound(self, name: str) -> TransformBound: folded = bound_of(self._stored_stages(name), rank).after(folded) return folded + def _pricing_bound(self, name: str) -> TransformBound: + """The map's bound as the PLAN prices it — headers and declarations, never a voxel. + + A field with no declared or recorded bound prices as zero displacement. The run never + trusts this window: a declared 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. + """ + if self.displacement is None or self.displacement.component_bound() is not None: + return self._bound(name) + rank = self._source_grid(name).rank + folded = TransformBound.exact(AffineMap.identity(rank)) + if self.transforms is not None: + folded = bound_of(self._stored_stages(name), rank).after(folded) + return folded + # ------------------------------------------------------------------ the contract def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: @@ -1548,13 +1581,20 @@ def _probe_cohort(self) -> str | None: "SimpleITK is not installed, and a stored transform is applied in physical space by" " it. Install it (pip install konfai[itk]) to stream this stage" ) - # The field's bound is the COHORT's, read from declarations and headers, so it is answered - # before any case has been seen -- and it is what a config-time probe is really asking. - if self.displacement is not None and self.displacement.component_bound() is None: - return self.displacement.undeclared_reason() + # The field group's HEADERS are the cohort's business here -- an unreadable entry anywhere + # under it fails both routes on whichever case reaches it. A field that merely records no + # bound streams: its windows are sized from the values the run reads (measured_region_source). + if self.displacement is not None: + self.displacement.component_bound() + if self.displacement.scan_failed: + return ( + "an entry in the field group could not be header-read, so what any region of it" + " must pull is unknown. Check the field store: one unreadable entry anywhere" + " under it falls the whole group back" + ) for name in self._grids: try: - self._bound(name) + self._pricing_bound(name) except TransformError as error: # Both halves of the refusal: the first says what is wrong, the second what to # change. A plan line carrying only the first tells the reader nothing to do. @@ -1572,7 +1612,34 @@ def stream_region_source( ) -> list[slice]: del source_spatial_shape, cache_attribute source, target = self._grids_of(name) - return list(source_window(target.sub_grid(tuple(target_slices)), source, self._bound(name))) + return list(source_window(target.sub_grid(tuple(target_slices)), source, self._pricing_bound(name))) + + @property + def measures_at_run(self) -> bool: + """Whether the run sizes this stage's windows from the data it reads. + + Only a field with NO declared or recorded bound: a bounded field keeps the declared + window, whose streamed result is bit-identical to the whole-volume path on a separable + map — measuring would tighten its windows at the price of that identity. Measuring is the + route for the field that could not stream at all before. + """ + return self.displacement is not None and self.displacement.component_bound() is None + + def measured_region_source( + self, name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + ) -> list[slice]: + """The region's source window, sized from the field itself — the read that samples also bounds. + + The field window a region needs is its own box, read for sampling regardless; the sup of + the values just read bounds every interpolated displacement in the region (a convex + combination cannot exceed the lattice values it blends), so the window is exact per region + — a quiet slab pays a quiet halo. ``max_displacement``, when declared, stays the cap those + values are checked against. + """ + del source_spatial_shape, cache_attribute + source, target = self._grids_of(name) + region = target.sub_grid(tuple(target_slices)) + return list(source_window(region, source, bound_of(self._stages(name, region), source.rank))) def stream_region( self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute @@ -1740,19 +1807,27 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut the mirror reason -- a question must not move the state a region read depends on. """ del group_dest + notes: list[str] = [] + if self.displacement is not None and self.displacement.component_bound() is None: + # Case-independent on purpose: the plan prints identical notes once, so this is one line + # for the stage rather than one per case. + notes.append( + "the field carries no bound, so each region's source window is sized from the field" + " values read at run; the read estimate prices the field as zero. Declare" + " max_displacement, or use a field with a recorded bound, to price exactly" + ) try: source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") - if missing & self._target.needs: - return None - covered = self._coverage(source, self._target.of(source, name), self._map_bound(name)) + if not missing & self._target.needs: + covered = self._coverage(source, self._target.of(source, name), self._map_bound(name)) + if covered < self._WORTH_SAYING: + notes.append( + f"case '{name}' covers {covered * 100:.1f}% of {self._target.describe()};" + f" the rest of what it writes is fill ({self.fill_value:g})" + ) except TransformError: - return None - if covered >= self._WORTH_SAYING: - return None - return ( - f"case '{name}' covers {covered * 100:.1f}% of {self._target.describe()};" - f" the rest of what it writes is fill ({self.fill_value:g})" - ) + pass + return "; ".join(notes) if notes else None # ------------------------------------------------------------------ the inverse @@ -2302,6 +2377,9 @@ def __init__( self.group = group #: The run's own roots, handed over by the owner; only consulted when there is no path. self.roots: list[Dataset] = [] + #: Whether the ``auto`` header scan DIED, as opposed to finding no bound: an unreadable + #: entry fails both routes at run, where a merely bound-less field streams (measured). + self.scan_failed = False self.auto = isinstance(max_displacement, str) and max_displacement.strip().lower() == "auto" if isinstance(max_displacement, str) and not self.auto: try: @@ -2350,6 +2428,7 @@ def component_bound(self) -> list[float] | None: 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 + self.scan_failed = True return self._auto_bound if bound and max(bound) > 0.0: self._auto_bound = bound diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 95c05684..72eba2e6 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -867,14 +867,15 @@ def test_a_field_beyond_its_declared_bound_is_refused(warped: tuple[Dataset, Dat 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.""" +def test_a_field_with_no_bound_still_streams(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """Warp's rule, and for the same reason: the run sizes each region's pull from the field + values it reads for sampling, so an undeclared bound is a pricing gap, not a fallback.""" 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 + stage = _warping(images, fields, max_displacement=0.0) + assert stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind is LocalityKind.REGRID + assert stage.measures_at_run + # A declared bound keeps the declared windows: their streamed result is pinned bit-identical. + assert not _warping(images, fields).measures_at_run # And without a field it is a region stage whatever the bound says. assert _stage_regrid_kind(images) is LocalityKind.REGRID diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index c5f5a798..7a871608 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -140,6 +140,9 @@ class _Case: transform: Transform group: str = "Intensity" atol: float = 0.0 + #: In the equivalence sweep. False for a case whose inputs this registry cannot build (a field + #: on disk): its streamed-equals-whole proof lives in its own test file, against real inputs. + sweep: bool = True # Only the transforms whose defaults are not a meaningful streaming case: a channel reduction needs a @@ -201,10 +204,10 @@ class _Case: # A stored map never factorises, so this is the grid_sample path (see _REGRID_ATOL). "ResampleTransform": [_Case(ResampleTransform({"transform": True}), atol=_REGRID_ATOL)], "Save": [_Case(Save("Dataset"))], - # Warp needs a field on disk to run, which this registry cannot build: with no declared - # displacement it declares WHOLE_VOLUME, so it stays out of the equivalence sweep below. Its - # streamed-equals-whole-volume proof lives in test_warp.py, where a field exists. - "Warp": [_Case(Warp(field="Dataset:h5", group="DVF"))], + # Warp needs a field on disk to run, which this registry cannot build. Its + # streamed-equals-whole-volume proof lives in test_warp.py, where a field exists — including + # the bound-less case, whose windows are measured from the field at run. + "Warp": [_Case(Warp(field="Dataset:h5", group="DVF"), sweep=False)], # Reduce is a cardinality marker the cohort engine splits out of the chain, never a per-case # stage: it declares WHOLE_VOLUME so a chain reaching the ordinary planner refuses rather than # streams, which is what puts it out of the equivalence sweep below. @@ -364,7 +367,10 @@ def _kind_of(case: _Case) -> LocalityKind: def _streamable_cases() -> list[_Case]: return [ - case for cls in _builtin_transforms() for case in _cases_of(cls) if _kind_of(case) not in _READ_REFUSED_KINDS + case + for cls in _builtin_transforms() + for case in _cases_of(cls) + if case.sweep and _kind_of(case) not in _READ_REFUSED_KINDS ] diff --git a/tests/unit/test_warp.py b/tests/unit/test_warp.py index b047c396..f385e1f7 100644 --- a/tests/unit/test_warp.py +++ b/tests/unit/test_warp.py @@ -181,14 +181,62 @@ def test_auto_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> assert locality.kind is LocalityKind.WHOLE_VOLUME -def test_auto_falls_back_to_the_whole_volume_when_no_field_recorded_a_bound(tmp_path: Path) -> None: - """Only an OME-Zarr field KonfAI wrote carries the bound, so `auto` must answer for the rest.""" - _source, _fields, _volume = _fixture(tmp_path) # an h5 field: no bound recorded +def test_a_field_with_no_bound_still_streams_with_windows_measured_at_run(tmp_path: Path) -> None: + """A bound-less field is not a whole-volume answer: the field window a region samples is read + for sampling regardless, and the sup of those very values sizes that region's source pull — + per region, so a quiet slab pays a quiet halo where the shifted one pays its shift.""" + _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 0.0, 0.0)) # 4 um along x alone + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto")) - locality = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto").patch_locality(_attributes()) + assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID - assert locality.kind is LocalityKind.WHOLE_VOLUME - assert locality.reason is not None and "no recorded bound" in locality.reason + target = (slice(4, 6), slice(4, 6), slice(4, 6)) + priced = warp.stream_region_source("CASE_000", target, [10, 12, 14], _attributes()) + measured = warp.measured_region_source("CASE_000", target, [10, 12, 14], _attributes()) + + # The plan prices as if the field were zero: the target's outer faces plus the taps' voxel. + assert [(part.start, part.stop) for part in priced] == [(2, 8), (2, 8), (2, 8)] + # The run pays the shift the values actually hold: 4 um at spacing 2 is 2 voxels, on x alone. + assert [(part.start, part.stop) for part in measured] == [(2, 8), (2, 8), (0, 10)] + + +def test_sizing_and_sampling_share_one_field_read(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The window that sizes a region's pull is the window the sampler needs next: one read.""" + _source, _fields, _volume = _fixture(tmp_path) + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto")) + displacement = warp.displacement + assert displacement is not None + reads: list[int] = [] + original = type(displacement).read + monkeypatch.setattr( + type(displacement), "read", lambda self, *args, **kwargs: (reads.append(1), original(self, *args, **kwargs))[1] + ) + + target = (slice(2, 5), slice(0, 12), slice(0, 14)) + warp.measured_region_source("CASE_000", target, [10, 12, 14], _attributes()) + _source_grid, target_grid = warp._grids_of("CASE_000") + warp._stages("CASE_000", target_grid.sub_grid(target)) + + assert len(reads) == 1 + + +def test_streamed_equals_whole_volume_with_no_declared_bound(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The measured windows carry the same claim as the declared ones: the same answer, region by + region, with several regions — and nothing was declared to make it true.""" + from konfai.data import patching as patching_module + + monkeypatch.setattr(patching_module, "_SWEEP_SLAB_ROWS", 3) + source, _fields, volume = _fixture(tmp_path, shift_um=(1.0, 2.0, 3.0)) + warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto") + + reference = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() + + manager = _manager(source, [warp, Save(f"{tmp_path / 'out'}:h5")]) + assert manager.can_stream_patch(0, apply_augmentations=False) + assert manager.materialize() is True + streamed, _ = Dataset(tmp_path / "out", "h5").read_data("CT", "CASE_000") + + np.testing.assert_allclose(streamed, reference, rtol=1e-5, atol=1e-4) def test_a_max_displacement_that_is_neither_a_number_nor_auto_is_refused() -> None: From 9c94f8935329769cdef1eaf8b47d79fe25cb5682 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 02:03:09 +0200 Subject: [PATCH 21/39] feat(data): drop max_displacement -- the field itself is the bound The parameter asked the user for a number the run now measures better: each region's source window is sized from the field values read for sampling. What remains needs no declaration -- the bound a STORE recorded at write time (KonfAI's OME-Zarr fields carry one) is read from headers, prices the plan exactly, and is checked per component against every region read: metadata that contradicts its data raises instead of sampling zeros. Removed from Resample, Warp and ResampleToReference; the impact_reg CLI loses --max-displacement. The streamed-equals-whole pin for a warped reference moves from bit-equality to the accepted band for maps that do not factorise: its exactness came from a declared bound large enough to make every window the whole source, and the measured windows are the region's own. The bit-exact claims stay on the separable path, which is window-independent by construction. --- apps/impact_reg/impact_reg_konfai/cli.py | 18 -- .../impact_reg_konfai/impact_reg.py | 20 +-- .../impact_reg_konfai/models/convexadam.py | 4 +- .../impact_reg_konfai/models/elastix.py | 4 +- .../models/elastix_engine.py | 5 +- .../impact_reg_konfai/models/fireants.py | 16 +- docs/source/config_guide/transform.md | 24 +-- .../source/reference/components/transforms.md | 6 +- konfai/data/transform.py | 162 ++++++------------ tests/unit/test_api.py | 4 +- tests/unit/test_resample.py | 5 - tests/unit/test_resample_to_reference.py | 54 +++--- .../unit/test_transform_locality_contract.py | 4 +- tests/unit/test_warp.py | 90 ++++------ .../test_write_pyramid_and_field_bound.py | 2 +- 15 files changed, 147 insertions(+), 271 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index 25b0778b..24d0b1f6 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -30,13 +30,6 @@ from impact_reg_konfai.impact_reg import ImpactRegKonfAIApp, get_available_presets -def _max_displacement(value: str) -> float | str: - """``auto`` or a distance in world units — the window bound a streamed field read needs.""" - if value.strip().lower() == "auto": - return "auto" - return float(value) - - def _paths(value: str) -> Path: return Path(value).resolve() @@ -142,16 +135,6 @@ def main() -> None: help="Write the displacement fields only: skip the moved image and Transform.h5, both derived " "from the field. For a caller that composes the field itself and would delete them.", ) - reg.add_argument( - "--max-displacement", - "--max_displacement", - dest="max_displacement", - type=_max_displacement, - default="auto", - help="Optional bound (world units) on the field: the moved image streams either way, each slab " - "sized from the field values it reads. A bound ('auto' reads the one OME-Zarr fields record) " - "lets the plan price the reads exactly, and a declared one is checked against every region read.", - ) _add_device(reg) _add_tmp_dir(reg) @@ -248,7 +231,6 @@ def _dispatch(args: argparse.Namespace, app: ImpactRegKonfAIApp, ev: argparse.Ar config_overrides=args.config_overrides, tmp_dir=args.tmp_dir, fields_only=args.fields_only, - max_displacement=args.max_displacement, ) elif args.command == "eval": diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 6ebf404d..5ff2d545 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -328,7 +328,6 @@ def register( config_overrides: list[str] | None = None, tmp_dir: Path | None = None, fields_only: bool = False, - max_displacement: float | str = "auto", ) -> None: """Register every case with the selected presets and ensemble their DVFs. @@ -346,11 +345,6 @@ def register( full-size rewrite -- worth it for a caller that wants a registration to look at, waste for one that composes the field with another and derives its own. A caller that reads only the fields should be able to say so rather than pay for outputs it deletes. - - ``max_displacement`` is optional: the moved image streams either way, each slab's source - window sized from the field values read for sampling. ``auto`` reads the bound a field - recorded (OME-Zarr fields carry one) so the plan prices the reads exactly; a distance in - world units declares that bound outright and is checked against every region actually read. """ # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs @@ -421,9 +415,7 @@ def register( # Every moved image in ONE streamed run: Resample adopts, per case, the grid of that # case's own DVF (a field is defined ON the fixed grid) and reads the field as the # map -- one interpolation, slab by slab, the whole cohort under one plan. - self._derive_moved( - dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet, max_displacement - ) + self._derive_moved(dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet) for case in cases: # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid field # as a SimpleITK transform. Inherently whole -- the .h5 format carries the full @@ -480,7 +472,6 @@ def _derive_moved( gpu: list[int], cpu: int | None, quiet: bool, - max_displacement: float | str, ) -> None: """The moved images, resampled from each moving through ITS displacement field — one run. @@ -488,7 +479,7 @@ def _derive_moved( moving, so deriving it belongs to this layer. THE GRID AND THE FORMAT FOLLOW THE FIELD, not the moving: a displacement field is defined ON the fixed grid, so ``reference: '{case}'`` adopts each case's own DVF grid, and the field is the map (``field_group``) — one - interpolation, streamed when the field's bound allows, whole-volume with the reason when not. + interpolation, streamed, each slab's source window sized from the field values it reads. """ from konfai.data.transform import Resample, Write @@ -506,12 +497,7 @@ def _derive_moved( { "Moving": { "Moved": [ - Resample( - reference="{case}", - reference_group="DVF", - field_group="DVF", - max_displacement=max_displacement, - ), + Resample(reference="{case}", reference_group="DVF", field_group="DVF"), Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), ] } diff --git a/apps/impact_reg/impact_reg_konfai/models/convexadam.py b/apps/impact_reg/impact_reg_konfai/models/convexadam.py index a80c4501..c41444ee 100644 --- a/apps/impact_reg/impact_reg_konfai/models/convexadam.py +++ b/apps/impact_reg/impact_reg_konfai/models/convexadam.py @@ -370,9 +370,7 @@ def _fine( # auto-hides it when stderr is not a TTY (e.g. under KonfAI/Slicer, where the outer "Prediction" bar # already reports progress), so captured logs stay clean; ``leave=False`` avoids stacking one bar per # patch. The observer is best-effort — if the filter emits no IterationEvent the bar just fills at the end. - progress = tqdm.tqdm( - total=self._iterations or None, desc="Registration", ncols=0, leave=False, disable=None - ) + progress = tqdm.tqdm(total=self._iterations or None, desc="Registration", ncols=0, leave=False, disable=None) def _update(*_: object) -> None: values = list(fine.GetMetricValuesPerIteration()) diff --git a/apps/impact_reg/impact_reg_konfai/models/elastix.py b/apps/impact_reg/impact_reg_konfai/models/elastix.py index 650eb6a5..3c5603c6 100644 --- a/apps/impact_reg/impact_reg_konfai/models/elastix.py +++ b/apps/impact_reg/impact_reg_konfai/models/elastix.py @@ -111,9 +111,7 @@ class ModelSpec: class ResolutionSpec: """One elastix resolution level: its iteration budget and the (self-configured) models compared there.""" - max_iterations: Annotated[ - int, Range(1, 100000), "Optimiser iterations spent at this resolution level." - ] + max_iterations: Annotated[int, Range(1, 100000), "Optimiser iterations spent at this resolution level."] models: dict[str, ModelSpec] diff --git a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py index af8f430c..60018640 100644 --- a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py +++ b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py @@ -278,7 +278,10 @@ def register( dst.symlink_to(model_path) args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)] - for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")): + for flag, mask, name in ( + ("-fMask", fixed_mask, "FixedMask.mha"), + ("-mMask", moving_mask, "MovingMask.mha"), + ): if _is_partial_mask(mask): mask_path = work / name sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path)) diff --git a/apps/impact_reg/impact_reg_konfai/models/fireants.py b/apps/impact_reg/impact_reg_konfai/models/fireants.py index 6e62a182..6220e7ec 100644 --- a/apps/impact_reg/impact_reg_konfai/models/fireants.py +++ b/apps/impact_reg/impact_reg_konfai/models/fireants.py @@ -485,9 +485,7 @@ def __init__( # Left to run this optimises nothing and returns the identity: a Moved equal to the moving # image and a zero field, which no downstream check tells apart from a pair that needed no # moving. - raise ValueError( - "linear_method='none' with deformable_method='none' leaves nothing to optimise." - ) + raise ValueError("linear_method='none' with deformable_method='none' leaves nothing to optimise.") self._deformable_metric = deformable_metric self._deformable_lr = float(deformable_lr) self._integrator_n = int(integrator_n) @@ -532,9 +530,7 @@ def com_phys(img: sitk.Image, mask: "sitk.Image | None") -> np.ndarray: # An all-zero (or all-negative) subject has no centre of mass to speak of; the frame # centre is the only defensible answer and matches what "cof" would have done. size = img.GetSize() - return np.asarray( - img.TransformContinuousIndexToPhysicalPoint([(extent - 1) / 2.0 for extent in size]) - ) + return np.asarray(img.TransformContinuousIndexToPhysicalPoint([(extent - 1) / 2.0 for extent in size])) index = np.array(np.nonzero(positive)) # (3, N) in z, y, x weight = array[positive] centre_voxel = (index * weight).sum(axis=1) / weight.sum() # z, y, x @@ -830,7 +826,9 @@ def __init__( Literal["mi", "cc", "mse"], "Similarity metric optimised during the affine (global) stage." ] = "mi", affine_lr: Annotated[ - float, Range(0.0, 10.0), "Gradient step size of the affine optimisation; higher converges faster but risks overshoot." + float, + Range(0.0, 10.0), + "Gradient step size of the affine optimisation; higher converges faster but risks overshoot.", ] = 0.003, moments_init: Annotated[ Literal["cof", "com"], @@ -855,9 +853,7 @@ def __init__( Literal["cc", "mi", "mse", "impact"], "Similarity metric for the deformable stage; 'impact' uses the IMPACT feature models under 'models'.", ] = "cc", - deformable_lr: Annotated[ - float, Range(0.0, 10.0), "Gradient step size of the deformable optimisation." - ] = 0.25, + deformable_lr: Annotated[float, Range(0.0, 10.0), "Gradient step size of the deformable optimisation."] = 0.25, integrator_n: Annotated[ int, Range(1, 100), diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index 758daaa1..e6317847 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -353,7 +353,6 @@ transforms: reference_group: CT field: ./Fields:mha field_group: DVF - max_displacement: 4.0 Write: {dataset: ./Registered:mha} ``` @@ -373,23 +372,24 @@ 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` is **optional**. The field window a region samples is its -own box, read for sampling regardless — and the sup of the values just read -bounds that region's source pull, so each slab pays exactly the halo *its* -displacements require, measured at run from a read the sampler needed anyway. -A declared bound (or `auto`, reading the one KonfAI records on a field it -writes) does two things: it lets the plan **price** the reads exactly — with -no bound the estimate assumes a zero field, and the plan says so — and it 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. +Nothing is declared about how far the field reaches. The field window a +region samples is its own box, read for sampling regardless — and the sup of +the values just read bounds that region's source pull, so each slab pays +exactly the halo *its* displacements require, measured at run from a read the +sampler needed anyway. A bound the **store recorded at write time** (KonfAI +records one on the OME-Zarr fields it writes) does two things without anyone +asking: it lets the plan **price** the reads exactly — without one the +estimate assumes a zero field, and the plan says so — and it is **checked +against every field region actually read**: a store whose data exceed its own +metadata raises rather than sampling zeros, which would show up as a dark rim +around the moved anatomy and nothing else. Naming no target grid is the shape update of an atlas build — the field applied on the case's *own* grid — and is the same stage with `reference` left out: ```yaml transforms: - Resample: {field: ./Fields:mha, field_group: DVF, max_displacement: 4.0} + Resample: {field: ./Fields:mha, field_group: DVF} Write: {dataset: ./Warped:mha} ``` diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 7975c39e..44206296 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -97,12 +97,12 @@ until it declares otherwise. | --- | --- | --- | --- | --- | --- | | `Padding` | `F.pad`; updates Origin. `mode` supports `"constant:"`. | `padding=[0,0,0,0,0,0], mode="constant", inverse=True` | **yes** | **yes** | no‡ | | `Crop` | Crop to foreground bounding box; caches the box; updates Origin. | `inverse=True` | **yes** | **yes** (pads back) | **yes** — once the `box` is on the case; the region is the patch translated | -| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region — and a declared/recorded `max_displacement` prices the plan exactly and is **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | +| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region — and a bound the store recorded at write time prices the plan exactly and is **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | | `ResampleToResolution` | Deprecated spelling of `Resample: {spacing: ...}`. | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** | | `ResampleToShape` | Deprecated spelling of `Resample: {shape: ...}`. | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** | -| `ResampleToReference` | Deprecated spelling of `Resample: {reference: ...}`. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** | **yes** | +| `ResampleToReference` | Deprecated spelling of `Resample: {reference: ...}`. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** | **yes** | | `ResampleTransform` | Deprecated spelling of `Resample: {transforms: ...}`. | `transforms` (required), `interpolation=None`, `fill=0.0`, `inverse=False` | no | no | **yes** | -| `Warp` | Deprecated spelling of `Resample: {field: ...}`. Note that `Resample` no longer requires the field and the case to share a grid. | `field` (required), `group=None`, `max_displacement=0.0`, `interpolation="linear"` | no | no | **yes** | +| `Warp` | Deprecated spelling of `Resample: {field: ...}`. Note that `Resample` no longer requires the field and the case to share a grid. | `field` (required), `group=None`, `interpolation="linear"` | no | no | **yes** | | `Canonical` | Reorient to canonical direction (3-D); updates Origin/Direction. | `inverse=True` | **yes** — a remap that transposes extents moves the patch grid | **yes** | **yes** — when the case's direction is a signed axis permutation; no on an oblique one (it is resampled) | | `Permute` | Permute spatial axes. `dims` is a pipe-separated axis list. | `dims="1\|0\|2", inverse=True` | **yes** | **yes** | **yes** — index remap | | `Flip` | Flip spatial axes. | `dims="1\|0\|2", inverse=True` | no | **yes** (self-inverse) | **yes** — index remap | diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 3186dd01..09f813d7 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -1204,9 +1204,9 @@ class Resample(TransformInverse): -- a theorem, not a sample of the boundary. A field on disk is read region by region, and the window a region samples is its own box: the sup of the values just read bounds that region's pull, so each slab pays exactly the halo ITS displacements require -- measured at run, from a - read the sampler needs regardless. ``max_displacement`` is optional: declared (or recorded by - the store at write time) it prices the plan exactly and is CHECKED against every region read; - absent, the plan prices the reads as if the field were zero and says so. + read the sampler needs regardless. Nothing is declared: a bound the STORE recorded at write + time (KonfAI's OME-Zarr fields carry one) prices the plan exactly and is CHECKED against every + region read; without one the plan prices the reads as if the field were zero, and says so. ``align`` decides where a ``spacing`` or a ``shape`` grid SITS, and it is the one silent choice in the family -- a quarter of a voxel of anatomy, made differently by every library that offers @@ -1227,12 +1227,12 @@ class Resample(TransformInverse): be ``fill`` from edge to edge, and an all-background member is a plausible, wrong contribution to a median. - A refusal the whole-volume path can serve -- an undeclared field bound, a case with no - geometry -- declares ``WHOLE_VOLUME`` with its reason and the run proceeds assembled: the chain - only stops being bounded, and says so in the plan. One that no route can serve -- a map that - cannot be decoded, read or inverted, or a disjoint case -- refuses as the plan is built, - before a byte is written. A case reaching only PART of the target grid is legal and common -- - the rest takes ``fill`` -- and the plan prints how much of the grid it covers. + A refusal the whole-volume path can serve -- a case with no geometry, an unreadable entry in + the field group -- declares ``WHOLE_VOLUME`` with its reason and the run proceeds assembled: + the chain only stops being bounded, and says so in the plan. One that no route can serve -- a + map that cannot be decoded, read or inverted, or a disjoint case -- refuses as the plan is + built, before a byte is written. A case reaching only PART of the target grid is legal and + common -- the rest takes ``fill`` -- and the plan prints how much of the grid it covers. """ def __init__( @@ -1245,7 +1245,6 @@ def __init__( transforms: dict[str, bool] | None = None, field: str | None = None, field_group: str | None = None, - max_displacement: float | str = 0.0, align: str = "extent", interpolation: str | None = None, fill: float = 0.0, @@ -1269,16 +1268,7 @@ def __init__( ) self.transforms = transforms 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"'Resample' 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: it sizes the" - " region a field is read from and means nothing without one.", - ) - self.displacement: _DisplacementSource | None = ( - _DisplacementSource(field, field_group, max_displacement) if declared else None - ) + self.displacement: _DisplacementSource | None = _DisplacementSource(field, field_group) if declared else None #: Per case: the grid its own header describes. Recorded where that header is in hand -- #: transform_shape, called for every case as the manager is built. A region read hands back #: the REGION's Origin, so a grid rebuilt from what a streamed region arrives with would @@ -1482,14 +1472,17 @@ def _stages(self, name: str, region: Grid) -> SpatialStages: return tuple(stages) def _bound(self, name: str) -> TransformBound: - """What the map is guaranteed to do — from declarations and coefficients, no voxel read.""" + """What the map is guaranteed to do — from recorded bounds and coefficients, no voxel read.""" rank = self._source_grid(name).rank folded = TransformBound.exact(AffineMap.identity(rank)) if self.displacement is not None: - declared = self.displacement.component_bound() - if declared is None: - raise TransformError(self.displacement.undeclared_reason()) - folded = TransformBound.shift(np.asarray(declared[:rank], dtype=np.float64)).after(folded) + recorded = self.displacement.component_bound() + if recorded is None: + raise TransformError( + "the field carries no recorded bound, so what it is guaranteed to do is unknown" + " before its values are read." + ) + folded = TransformBound.shift(np.asarray(recorded[:rank], dtype=np.float64)).after(folded) if self.transforms is not None: folded = bound_of(self._stored_stages(name), rank).after(folded) return folded @@ -1524,7 +1517,7 @@ def transform_shape(self, group_src: str, name: str, shape: list[int], cache_att def _require_runnable(self, name: str) -> None: """Refuse AT PLAN TIME a map neither route can apply. - A refusal the whole-volume path can serve — an undeclared field bound — stays a locality + A refusal the whole-volume path can serve — a case with no geometry — stays a locality answer, and the run proceeds assembled. A stored transform that cannot be decoded, read or inverted fails the streamed path and the whole-volume one at the same line, so declaring WHOLE_VOLUME for it would print a plan the run then contradicts by dying per case, after @@ -1633,8 +1626,8 @@ def measured_region_source( The field window a region needs is its own box, read for sampling regardless; the sup of the values just read bounds every interpolated displacement in the region (a convex combination cannot exceed the lattice values it blends), so the window is exact per region - — a quiet slab pays a quiet halo. ``max_displacement``, when declared, stays the cap those - values are checked against. + — a quiet slab pays a quiet halo. A bound the store recorded stays the cap those values + are checked against. """ del source_spatial_shape, cache_attribute source, target = self._grids_of(name) @@ -1813,8 +1806,8 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut # for the stage rather than one per case. notes.append( "the field carries no bound, so each region's source window is sized from the field" - " values read at run; the read estimate prices the field as zero. Declare" - " max_displacement, or use a field with a recorded bound, to price exactly" + " values read at run; the read estimate prices the field as zero. A field with a" + " recorded bound (KonfAI records one on the OME-Zarr fields it writes) prices exactly" ) try: source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") @@ -1967,7 +1960,6 @@ def __init__( dataset: str | None = None, field: str | None = None, field_group: str | None = None, - max_displacement: float | str = 0.0, fill: float = 0.0, interpolation: str | None = None, inverse: bool = True, @@ -1983,7 +1975,6 @@ def __init__( reference_dataset=dataset, field=field, field_group=field_group, - max_displacement=max_displacement, fill=fill, interpolation=interpolation, inverse=inverse, @@ -2355,12 +2346,7 @@ class _DisplacementSource: speaks as ``Resample`` and names ``field_group``, the argument the user declared. """ - def __init__( - self, - field: str | None, - group: str | None, - max_displacement: float | str, - ) -> None: + def __init__(self, field: str | None, group: str | None) -> None: # 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. @@ -2377,39 +2363,24 @@ def __init__( self.group = group #: The run's own roots, handed over by the owner; only consulted when there is no path. self.roots: list[Dataset] = [] - #: Whether the ``auto`` header scan DIED, as opposed to finding no bound: an unreadable + #: Whether the header scan DIED, as opposed to finding no recorded bound: an unreadable #: entry fails both routes at run, where a merely bound-less field streams (measured). self.scan_failed = False - 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"'Resample' 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 else float(max_displacement) - self._auto_bound: list[float] | None = None - self._auto_resolved = False + self._recorded_bound: list[float] | None = None + self._scan_resolved = False def component_bound(self) -> list[float] | None: - """The per-component bound this stage warps within, or ``None`` when it has none. + """The per-component bound the STORE recorded, or ``None`` when it recorded none. - For ``auto``, the largest bound any field in the group recorded, read from headers alone and - memoized. If a single entry carries no bound the answer is ``None``: a maximum over the - 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. + The largest bound any field in the group recorded at write time, read from headers alone + and memoized — nobody declares anything. If a single entry carries no bound the answer is + ``None``: a maximum over the others would be a bound for them and a guess for that one. + Per component, in the field's own (x, y, z) order, because these grids are anisotropic: + one collapsed maximum over-reads the fine axes. """ - 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 - self._auto_resolved = True + if self._scan_resolved: + return self._recorded_bound + self._scan_resolved = True from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE bound: list[float] = [] @@ -2424,29 +2395,15 @@ def component_bound(self) -> list[float] | None: 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 + return self._recorded_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 self.scan_failed = True - return self._auto_bound + return self._recorded_bound if bound and max(bound) > 0.0: - self._auto_bound = bound - return self._auto_bound - - 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" - ) + self._recorded_bound = bound + return self._recorded_bound def group_for(self, name: str | None) -> str: if self.group is not None: @@ -2501,7 +2458,7 @@ def read(self, name: str, region: tuple[slice, ...] | None, channels: int) -> to 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. + """The store's recorded bound is a promise about the region 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. @@ -2510,47 +2467,28 @@ def check_bound(self, field: torch.Tensor, name: str) -> None: 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) + recorded = bound[component] if component < len(bound) else max(bound) largest = float(field[component].abs().max()) - if largest > declared: + if largest > recorded: raise TransformError( f"The field for case '{name}' displaces up to {largest:.3f} on component" - f" {component}, beyond the {declared:.3f} 'Resample' 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.", + f" {component}, beyond the {recorded:.3f} its store recorded — the bound" + " 'Resample' sized its region from.", + "The store's metadata contradicts its data: rewrite the field so the recorded" + " bound holds, or strip the stale bound so the windows are measured instead.", ) -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 - - class Warp(Resample): """Deprecated spelling of ``Resample: {field: ...}`` — a warp on the case's own grid.""" - def __init__( - self, - field: str, - group: str | None = None, - max_displacement: float | str = 0.0, - interpolation: str = "linear", - ) -> None: + def __init__(self, field: str, group: str | None = None, interpolation: str = "linear") -> None: if not field or not str(field).strip(): raise TransformError( "'Warp' needs a 'field': the displacement field to resample through.", - "Declare it, e.g. Resample: {field: ./DVF:omezarr, max_displacement: 250.0}.", + "Declare it, e.g. Resample: {field: ./DVF:omezarr}.", ) - super().__init__( - field=field, - field_group=group, - max_displacement=max_displacement, - interpolation=interpolation, - inverse=False, - ) + super().__init__(field=field, field_group=group, interpolation=interpolation, inverse=False) class Reduce(Transform): diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py index c213b432..3842d5c4 100644 --- a/tests/unit/test_api.py +++ b/tests/unit/test_api.py @@ -66,8 +66,8 @@ def test_a_repeated_mapping_stage_is_qualified_by_resolution() -> None: def test_a_subclass_delegating_to_super_keeps_its_own_spelling() -> None: """``Warp(field=...)`` expands into ``Resample`` arguments internally; the recorded spelling is the caller's, so the tree references ``Warp`` with the caller's kwargs and rebinds identically.""" - stage = Warp(field="./DVF:omezarr", max_displacement=120.0) - assert stage._konfai_given == {"field": "./DVF:omezarr", "max_displacement": 120.0} + stage = Warp(field="./DVF:omezarr", group="DVF") + assert stage._konfai_given == {"field": "./DVF:omezarr", "group": "DVF"} def test_the_chain_tree_is_the_yaml_subtree() -> None: diff --git a/tests/unit/test_resample.py b/tests/unit/test_resample.py index 80ccb135..9f1225a0 100644 --- a/tests/unit/test_resample.py +++ b/tests/unit/test_resample.py @@ -287,11 +287,6 @@ def test_a_change_of_grid_and_a_warp_are_one_interpolation(tmp_path) -> None: np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-3) -def test_a_bound_declared_without_a_field_is_refused() -> None: - with pytest.raises(TransformError, match="no field to apply"): - Resample(spacing=[1.0, 1.0, 1.0], max_displacement=10.0) - - # ------------------------------------------------------------------ which loop runs diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 72eba2e6..14f68050 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -546,7 +546,6 @@ def test_the_inverse_says_why_when_the_forward_left_no_stack(dataset: Dataset) - # 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: @@ -584,7 +583,6 @@ def _warping(images: Dataset, fields: Dataset, **kwargs: object) -> ResampleToRe "group": "Reference", "field": f"{fields.filename}:h5", "field_group": "DVF", - "max_displacement": _BOUND, "fill": _FILL, **kwargs, } @@ -741,7 +739,6 @@ def test_the_field_components_are_not_reversed(tmp_path: Path) -> None: group="Reference", field=f"{fields.filename}:h5", field_group="DVF", - max_displacement=4.0, fill=0.0, ) stage.set_datasets([images]) @@ -797,7 +794,13 @@ def test_it_interpolates_once_not_twice(warped: tuple[Dataset, Dataset, np.ndarr 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.""" + """Region by region through a field, against the same chain run whole. + + Within the accepted band for a map that does not factorise, not bit-for-bit: the blend goes + through ``grid_sample``, whose float arithmetic depends on the window extent, and the measured + windows are the region's own rather than the whole volume. The bit-exact claims live on the + separable path, which is window-independent by construction. + """ images, fields, volume = warped def manager() -> DatasetManager: @@ -824,7 +827,7 @@ def manager() -> DatasetManager: 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) + np.testing.assert_allclose(got, whole, rtol=1e-5, atol=1e-4) def test_the_warped_run_never_assembles_the_volume(warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path) -> None: @@ -855,28 +858,35 @@ def refuse(*args: object, **kwargs: object) -> None: 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. +def test_a_field_beyond_its_recorded_bound_is_refused( + warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path +) -> None: + """The store's bound is a promise about what was read; data that break 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) + images, _fields, volume = warped + attributes = _attributes(_FIELD_ORIGIN, _FIELD_SPACING) + attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.asarray([0.5, 0.5, 0.5]) + lying = Dataset(tmp_path / "lying", "mha") + lying.write("DVF", _CASE, _displacement(), attributes) + stage = ResampleToReference( + entry=_CASE, group="Reference", field=f"{tmp_path / 'lying'}:mha", field_group="DVF", fill=_FILL + ) + stage.set_datasets([images]) 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_still_streams(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: """Warp's rule, and for the same reason: the run sizes each region's pull from the field - values it reads for sampling, so an undeclared bound is a pricing gap, not a fallback.""" + values it reads for sampling, so a missing recorded bound is a pricing gap, not a fallback.""" images, fields, _volume = warped - stage = _warping(images, fields, max_displacement=0.0) + stage = _warping(images, fields) assert stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind is LocalityKind.REGRID assert stage.measures_at_run - # A declared bound keeps the declared windows: their streamed result is pinned bit-identical. - assert not _warping(images, fields).measures_at_run - # And without a field it is a region stage whatever the bound says. + # And without a field it is a region stage, and nothing is measured. assert _stage_regrid_kind(images) is LocalityKind.REGRID @@ -918,7 +928,7 @@ def test_fields_can_live_beside_the_cases(warped: tuple[Dataset, Dataset, np.nda """ 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 = ResampleToReference(entry=_CASE, group="Reference", field_group="DVF", fill=_FILL) beside.set_datasets([images]) got = beside(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) @@ -928,7 +938,7 @@ 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: +def test_the_recorded_bound_is_read_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 @@ -942,23 +952,13 @@ def test_auto_reads_the_bound_from_every_root_not_the_first(tmp_path: Path) -> N 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 = ResampleToReference(entry="CASE_1", group="Reference", field_group="DVF") 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. - - 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) - - def test_the_two_gathers_obey_the_same_rules_through_an_identity_field(tmp_path: Path) -> None: """One arithmetic, two loops: per-axis maps, and eight corners at a coordinate volume. diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index 7a871608..bd317756 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -195,9 +195,7 @@ class _Case: # normalises by the extent it is handed, and a patch is handed a window. That is the one # place a streamed answer is not bit-identical to the whole-volume one, and the atol says so. _Case( - ResampleToReference( - entry=_CASE_NAME, group="Reference", field_group="Field", max_displacement=_FIELD_BOUND - ), + ResampleToReference(entry=_CASE_NAME, group="Reference", field_group="Field"), atol=_REGRID_ATOL, ), ], diff --git a/tests/unit/test_warp.py b/tests/unit/test_warp.py index f385e1f7..ddb48af4 100644 --- a/tests/unit/test_warp.py +++ b/tests/unit/test_warp.py @@ -17,7 +17,8 @@ """``Warp`` resamples a case through a displacement field, region by region. The claim under test is the one that matters for a volume larger than memory: the streamed result -equals the whole-volume one, and the declared displacement bound is verified rather than trusted.""" +equals the whole-volume one, each region's window is sized from the field values it reads, and a +bound the store recorded is verified rather than trusted.""" from pathlib import Path @@ -28,7 +29,7 @@ from konfai.data.transform import LocalityKind, RegionContext, Save, Warp from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, Attribute, Dataset from konfai.utils.errors import TransformError -from konfai.utils.ome_zarr import _zarr_v3_available +from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE, _zarr_v3_available pytest.importorskip("SimpleITK") @@ -87,22 +88,23 @@ def _recorded(warp: Warp, attribute: Attribute | None = None, shape: tuple[int, return warp -def test_the_source_region_is_the_target_grown_by_the_declared_displacement() -> None: - """A warp is a regrid onto the case's own grid, and its window is the bound in voxels. +def test_the_source_region_is_the_target_grown_by_the_field_reach(tmp_path: Path) -> None: + """A warp is a regrid onto the case's own grid, and its window is the field's reach in voxels. Spacing is (x=2, y=1, z=1), so in array order (z, y, x) 4 um of displacement is 4, 4 and 2 voxels -- plus the one voxel the linear taps reach. Declared as REGRID and not HALO because the window is derived from the case's GEOMETRY: see the oblique case below, which a per-axis halo cannot express at all. """ - warp = _recorded(Warp(field="./x:h5", group="DVF", max_displacement=4.0)) + _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 4.0, 4.0)) + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID target = (slice(4, 6), slice(4, 6), slice(4, 6)) - window = warp.stream_region_source("CASE_000", target, [10, 12, 14], _attributes()) + window = warp.measured_region_source("CASE_000", target, [10, 12, 14], _attributes()) # The rule, written out: the region's OUTER faces (start - 0.5 .. stop - 0.5) in world units, - # grown by the declared 4 um, back to indices, floor/ceil, one voxel of margin for the taps. + # grown by the field's 4 um, back to indices, floor/ceil, one voxel of margin for the taps. extents, per_voxel = (10, 12, 14), (1.0, 1.0, 2.0) # array order (z, y, x) expected = [] for axis, extent in enumerate(extents): @@ -112,7 +114,7 @@ def test_the_source_region_is_the_target_grown_by_the_declared_displacement() -> assert [(part.start, part.stop) for part in window] == expected -def test_an_oblique_case_grows_its_window_on_every_axis() -> None: +def test_an_oblique_case_grows_its_window_on_every_axis(tmp_path: Path) -> None: """The bug a per-axis halo hid: a displacement along x reaches into y and z when the axes turn. ``Warp`` used to convert a world bound to a halo per ARRAY axis, which silently assumed the @@ -124,22 +126,20 @@ def test_an_oblique_case_grows_its_window_on_every_axis() -> None: cos, sin = float(np.cos(angle)), float(np.sin(angle)) turned["Direction"] = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]).reshape(-1) - warp = _recorded(Warp(field="./x:h5", group="DVF", max_displacement=4.0), turned) + _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 0.0, 0.0)) + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF"), turned) target = (slice(5, 6), slice(5, 6), slice(5, 6)) - window = warp.stream_region_source("CASE_000", target, [10, 12, 14], turned) + window = warp.measured_region_source("CASE_000", target, [10, 12, 14], turned) widths = [part.stop - part.start for part in window] assert all(width > 1 for width in widths), f"a turned case reaches on every axis, got {widths}" @_needs_rfc5 -def test_auto_reads_the_bound_the_fields_recorded_when_they_were_written(tmp_path: Path) -> None: - """``max_displacement: auto`` is the number the producer already knew. - - A field records its own per-component bound at write time, so asking the user to measure it is - asking for something the store can answer from a header. Per component and not one collapsed - maximum: these grids are anisotropic, and one number over-reads the fine axes. - """ +def test_the_bound_the_fields_recorded_prices_the_plan(tmp_path: Path) -> None: + """The recorded bound is the number the producer already knew — read from headers, declared by + nobody. It sizes the plan's (headers-only) windows; per component and not one collapsed + maximum, because these grids are anisotropic and one number over-reads the fine axes.""" fields = Dataset(tmp_path / "dvf", "omezarr") for case, shift in (("CASE_000", (1.0, 2.0, 3.0)), ("CASE_001", (0.5, 6.0, 1.0))): field = np.zeros((3, 4, 5, 6), dtype=np.float32) @@ -149,7 +149,7 @@ def test_auto_reads_the_bound_the_fields_recorded_when_they_were_written(tmp_pat attribute[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" fields.write("DVF", case, field, attribute) - warp = Warp(field=f"{tmp_path / 'dvf'}:omezarr", group="DVF", max_displacement="auto") + warp = Warp(field=f"{tmp_path / 'dvf'}:omezarr", group="DVF") locality = warp.patch_locality(_attributes()) # The cohort's bound is (x=1.0, y=6.0, z=3.0); spacing in array order (z, y, x) is (1, 1, 2), so @@ -161,7 +161,7 @@ def test_auto_reads_the_bound_the_fields_recorded_when_they_were_written(tmp_pat assert [part.start for part in window] == starts -def test_auto_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> None: +def test_the_header_scan_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> None: """The whole group is header-read, including entries this run never warps. A directory store lists its entries from the filesystem alone, so a corrupt one is only met at @@ -176,7 +176,7 @@ def test_auto_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> # The first entry the scan meets, so the read reaches it before any bound-less entry ends the scan. (tmp_path / "dvf" / "CASE_000" / "DVF.mha").write_bytes(b"not an image") - locality = Warp(field=f"{tmp_path / 'dvf'}:mha", group="DVF", max_displacement="auto").patch_locality(_attributes()) + locality = Warp(field=f"{tmp_path / 'dvf'}:mha", group="DVF").patch_locality(_attributes()) assert locality.kind is LocalityKind.WHOLE_VOLUME @@ -186,7 +186,7 @@ def test_a_field_with_no_bound_still_streams_with_windows_measured_at_run(tmp_pa for sampling regardless, and the sup of those very values sizes that region's source pull — per region, so a quiet slab pays a quiet halo where the shifted one pays its shift.""" _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 0.0, 0.0)) # 4 um along x alone - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto")) + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID @@ -203,7 +203,7 @@ def test_a_field_with_no_bound_still_streams_with_windows_measured_at_run(tmp_pa def test_sizing_and_sampling_share_one_field_read(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """The window that sizes a region's pull is the window the sampler needs next: one read.""" _source, _fields, _volume = _fixture(tmp_path) - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto")) + warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) displacement = warp.displacement assert displacement is not None reads: list[int] = [] @@ -220,33 +220,9 @@ def test_sizing_and_sampling_share_one_field_read(tmp_path: Path, monkeypatch: p assert len(reads) == 1 -def test_streamed_equals_whole_volume_with_no_declared_bound(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - """The measured windows carry the same claim as the declared ones: the same answer, region by - region, with several regions — and nothing was declared to make it true.""" - from konfai.data import patching as patching_module - - monkeypatch.setattr(patching_module, "_SWEEP_SLAB_ROWS", 3) - source, _fields, volume = _fixture(tmp_path, shift_um=(1.0, 2.0, 3.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement="auto") - - reference = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() - - manager = _manager(source, [warp, Save(f"{tmp_path / 'out'}:h5")]) - assert manager.can_stream_patch(0, apply_augmentations=False) - assert manager.materialize() is True - streamed, _ = Dataset(tmp_path / "out", "h5").read_data("CT", "CASE_000") - - np.testing.assert_allclose(streamed, reference, rtol=1e-5, atol=1e-4) - - -def test_a_max_displacement_that_is_neither_a_number_nor_auto_is_refused() -> None: - with pytest.raises(TransformError, match="neither a number nor 'auto'"): - Warp(field="./x:h5", group="DVF", max_displacement="lots") - - def test_a_constant_shift_moves_the_volume_by_that_many_voxels(tmp_path: Path) -> None: _source, _fields, volume = _fixture(tmp_path, shift_um=(0.0, 0.0, 3.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement=3.0) + warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") moved = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() @@ -255,12 +231,13 @@ def test_a_constant_shift_moves_the_volume_by_that_many_voxels(tmp_path: Path) - def test_streamed_equals_whole_volume(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - """The claim that matters: the same answer, region by region, with several regions.""" + """The claim that matters: the same answer, region by region, with several regions — and + nothing was declared to make it true, the windows being measured from the field itself.""" from konfai.data import patching as patching_module monkeypatch.setattr(patching_module, "_SWEEP_SLAB_ROWS", 3) source, _fields, volume = _fixture(tmp_path, shift_um=(1.0, 2.0, 3.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement=4.0) + warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") reference = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() @@ -272,15 +249,20 @@ def test_streamed_equals_whole_volume(tmp_path: Path, monkeypatch: pytest.Monkey np.testing.assert_allclose(streamed, reference, rtol=1e-5, atol=1e-4) -def test_a_field_beyond_the_declared_bound_raises(tmp_path: Path) -> None: - """Declared, then verified: sampling outside what was read would show as a dark rim and nothing - else, so the mismatch is raised instead. +def test_a_field_beyond_its_recorded_bound_raises(tmp_path: Path) -> None: + """The store's metadata is a promise about what was read; a store whose data break it must not + sample zeros — a dark rim around the moved anatomy and nothing else to see. Checked per component, the way the halo is derived: a field under the collapsed maximum can still exceed the bound on one axis, and that axis is the one whose halo was too small. """ _source, _fields, volume = _fixture(tmp_path, shift_um=(0.0, 0.0, 9.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement=1.0) + attributes = _attributes() + attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.asarray([1.0, 1.0, 1.0]) + field = np.zeros((3, 10, 12, 14), dtype=np.float32) + field[2] = 9.0 + Dataset(tmp_path / "lying", "mha").write("DVF", "CASE_000", field, attributes) + warp = Warp(field=f"{tmp_path / 'lying'}:mha", group="DVF") _recorded(warp) with pytest.raises(TransformError, match=r"on component 2, beyond the 1\.000"): @@ -297,7 +279,7 @@ def test_a_field_with_the_wrong_component_count_is_named(tmp_path: Path) -> None rng = np.random.default_rng(1) Dataset(tmp_path / "src", "h5").write("CT", "CASE_000", rng.random((1, 4, 4, 4)).astype(np.float32), _attributes()) Dataset(tmp_path / "dvf", "h5").write("DVF", "CASE_000", np.zeros((2, 4, 4, 4), np.float32), _attributes()) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF", max_displacement=1.0) + warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") with pytest.raises(TransformError, match="component"): warp("CASE_000", torch.zeros(1, 4, 4, 4), _attributes()) diff --git a/tests/unit/test_write_pyramid_and_field_bound.py b/tests/unit/test_write_pyramid_and_field_bound.py index 481761fe..d1f6a6a5 100644 --- a/tests/unit/test_write_pyramid_and_field_bound.py +++ b/tests/unit/test_write_pyramid_and_field_bound.py @@ -216,7 +216,7 @@ def test_the_recorded_bound_reaches_each_axis_by_its_own_spacing_under_anisotrop field[0, 4, 4, 4], field[1, 2, 2, 2], field[2, 1, 1, 1] = 917.5, -640.25, 96.0 store = tmp_field_store(field) - warp = Warp(field=f"{store}:omezarr", group="DVF", max_displacement="auto") + warp = Warp(field=f"{store}:omezarr", group="DVF") attribute = Attribute() attribute["Spacing"] = np.array([30.08, 30.08, 40.0]) # stored (x, y, z) attribute["Origin"] = np.zeros(3) From 7a809b3b14a5ce3d555fc48755d502294a116c6f Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 02:31:44 +0200 Subject: [PATCH 22/39] feat(data)!: one Resample, no other spelling ResampleToResolution, ResampleToShape, ResampleToReference, ResampleTransform and Warp are removed: five ways to reference the one stage, kept as forwarding shells. A config names Resample and says which grid (spacing, shape, reference) and which map (field, transforms); two Resample stages in one chain spell the second module-qualified (konfai.data.transform:Resample), which the strict grammar and the loader already resolve. A blank reference is refused at construction. The locality-contract registry folds the family's cases under one Resample entry; the docs lose the deprecated-spelling rows and notes. BREAKING CHANGE: published configs using the old names must rename the key -- ResampleToResolution: {spacing: [...]} becomes Resample: {spacing: [...]}. The HF bundle configs move with the konfai version they pin. --- docs/source/concepts/streaming.md | 4 +- docs/source/config_guide/prediction.md | 2 +- docs/source/config_guide/transform.md | 14 +- docs/source/examples/visual-gallery.md | 6 +- .../source/reference/components/transforms.md | 5 - docs/source/troubleshooting.md | 2 +- docs/source/usage/large-images.md | 6 +- konfai/data/transform.py | 82 +---- konfai/transformer.py | 2 +- .../test_konfai_streamed_prediction.py | 8 +- tests/integration/test_transform_example.py | 2 +- tests/unit/test_api.py | 15 +- tests/unit/test_case_expansion.py | 3 +- tests/unit/test_itk_transforms.py | 6 +- tests/unit/test_resample.py | 310 +----------------- tests/unit/test_resample_to_reference.py | 54 +-- tests/unit/test_resample_transform.py | 32 +- tests/unit/test_streamed_read_dispatcher.py | 17 +- tests/unit/test_streamed_write_dispatcher.py | 16 +- tests/unit/test_transform.py | 27 +- .../unit/test_transform_locality_contract.py | 51 ++- tests/unit/test_transformer_workflow.py | 15 +- tests/unit/test_warp.py | 40 +-- .../test_write_pyramid_and_field_bound.py | 4 +- 24 files changed, 160 insertions(+), 563 deletions(-) diff --git a/docs/source/concepts/streaming.md b/docs/source/concepts/streaming.md index 2d0ef30d..f03a3cfa 100644 --- a/docs/source/concepts/streaming.md +++ b/docs/source/concepts/streaming.md @@ -163,7 +163,7 @@ before it, and so on down to one region on disk. KonfAI reads that one region and runs the chain forward over it. `[Dilate(1), Gradient()]` is two halos that add, `[Canonical(), Permute('2|1|0')]` is two reorientations that pull through each other — both stream from a single bounded read. The reorient-resample-pad -stack `[Canonical, ResampleToResolution, Padding]` streams on the **write** side but +stack `[Canonical, Resample, Padding]` streams on the **write** side but not on the read side: `Padding` declares no forward locality, so it inherits `WHOLE_VOLUME` and the chain is refused at that stage. The one thing a region stage cannot do is read a statistic of its own input, because @@ -184,7 +184,7 @@ streams and `Dilate(5)` does not. | `HALO` | `Dilate(n>0)`, `Gradient` | | `ORIENTATION` | `Flip`, `Permute`, `Canonical` (only on axis-aligned direction cosines) | | `CROP` | `Crop` (only once its box is on the case) | -| `REGRID` | `Resample` — and its spellings `ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, `ResampleTransform`, `Warp` | +| `REGRID` | `Resample` — whichever grid and map it is given | Augmentations declare per **(case, draw)** — two copies of the same case can answer differently. `Permute`, `Flip` (when `vector_field: false`), and `Rotate` diff --git a/docs/source/config_guide/prediction.md b/docs/source/config_guide/prediction.md index 85381430..2c64862e 100644 --- a/docs/source/config_guide/prediction.md +++ b/docs/source/config_guide/prediction.md @@ -162,7 +162,7 @@ the exported output name stays consistent. identically to the assembled volume (a single augmentation, a voxel-local reduction, and an `mha`/`h5`/`omezarr` destination), each slab is written to disk as its patches complete, bounding RAM at one patch window instead of the whole volume. Geometry inverses stream too, composed in any number -(`Canonical`/`Flip`/`Permute`, `Padding`, nearest-mode `ResampleToResolution`/`ResampleToShape`): +(`Canonical`/`Flip`/`Permute`, `Padding`, a nearest-mode `Resample`): each slab is remapped, cropped, or resampled through a sliding window straight to its written region. A chain streaming cannot honour streams its pointwise prefix into a light post-reduction buffer and runs the rest whole-volume on it. Streamed outputs match the assembled path voxel for voxel on a given diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index e6317847..76c9aeba 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -302,13 +302,6 @@ outer faces coincide — while `origin` keeps voxel zero's centre where it is. A quarter of a voxel of anatomy separates them, and a `reference` states its own placement and ignores this. -```{note} -`ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, -`ResampleTransform` and `Warp` are still accepted, and are now thin spellings of -this one stage: `Resample: {spacing: …}`, `{shape: …}`, `{reference: …}`, -`{transforms: …}` and `{field: …}` respectively. -``` - ### `Resample: {reference: …}`: making `strict` true rather than waived A cohort as acquired rarely passes `strict`: extents differ, and origins can @@ -394,10 +387,9 @@ transforms: ``` ```{note} -This was `Warp`, which required the field and the case to share a grid. They no -longer have to: the field is read at each target voxel's world position on the -field's own grid, so a field solved at 120 µm moves a volume stored at 30 µm -without being upsampled first. +The field and the case need not share a grid: the field is read at each target +voxel's world position on the field's own grid, so a field solved at 120 µm +moves a volume stored at 30 µm without being upsampled first. ``` Naming an image rather than fifteen numbers is deliberate. A grid is an extent diff --git a/docs/source/examples/visual-gallery.md b/docs/source/examples/visual-gallery.md index 59da2c55..517e3ea1 100644 --- a/docs/source/examples/visual-gallery.md +++ b/docs/source/examples/visual-gallery.md @@ -64,8 +64,8 @@ operation when predictions are written.
  • Clip output in the same display window, with intensities above 100 HU flattened so bright bone disappears.
    ClipValues above 100 HU flattened · same display windowMIN −1000.00 · MEAN −575.52 · MAX 100.00
  • Normalize output mapped to the numerical range minus one to one.
    NormalizeTarget numerical range [−1, 1]MIN −1.00 · MEAN −0.23 · MAX 1.00
  • Standardize output centered at zero and scaled by its standard deviation.
    StandardizeAutomatic mean and standard deviationMIN −0.87 · MEAN 0.00 · MAX 3.11
    • -
  • ResampleToShape output on a 220 by 220 grid.
    Resample to shape219 × 222 → 220 × 220 voxelsOUTPUT GRID 220 × 220
    • -
  • ResampleToResolution output showing the anisotropic target spacing.
    Resample spacing0.80 × 0.80 → 1.25 × 0.55 mmGEOMETRY-AWARE INTERPOLATION
    • +
  • Resample output on a 220 by 220 grid.
    Resample to shape219 × 222 → 220 × 220 voxelsOUTPUT GRID 220 × 220
    • +
  • Resample output showing the anisotropic target spacing.
    Resample spacing0.80 × 0.80 → 1.25 × 0.55 mmGEOMETRY-AWARE INTERPOLATION
  • Padding output with an expanded image grid.
    PaddingConstant −1 border around the source grid[28, 28, 18, 18] VOXELS
  • Crop output restricted to the configured central region.
    CropConfigured central ROI with geometry updateROI 177 × 178 VOXELS
  • Permute output after swapping the two spatial axes.
    Permute axesSwap the two spatial dimensions219 × 222 → 222 × 219
    • @@ -87,7 +87,7 @@ groups_dest: min_value: -1 max_value: 1 inverse: true - ResampleToResolution: + Resample: spacing: [1.25, 0.55] inverse: true Padding: diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 44206296..03b87897 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -98,11 +98,6 @@ until it declares otherwise. | `Padding` | `F.pad`; updates Origin. `mode` supports `"constant:"`. | `padding=[0,0,0,0,0,0], mode="constant", inverse=True` | **yes** | **yes** | no‡ | | `Crop` | Crop to foreground bounding box; caches the box; updates Origin. | `inverse=True` | **yes** | **yes** (pads back) | **yes** — once the `box` is on the case; the region is the patch translated | | `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region — and a bound the store recorded at write time prices the plan exactly and is **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | -| `ResampleToResolution` | Deprecated spelling of `Resample: {spacing: ...}`. | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** | -| `ResampleToShape` | Deprecated spelling of `Resample: {shape: ...}`. | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** | -| `ResampleToReference` | Deprecated spelling of `Resample: {reference: ...}`. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** | **yes** | -| `ResampleTransform` | Deprecated spelling of `Resample: {transforms: ...}`. | `transforms` (required), `interpolation=None`, `fill=0.0`, `inverse=False` | no | no | **yes** | -| `Warp` | Deprecated spelling of `Resample: {field: ...}`. Note that `Resample` no longer requires the field and the case to share a grid. | `field` (required), `group=None`, `interpolation="linear"` | no | no | **yes** | | `Canonical` | Reorient to canonical direction (3-D); updates Origin/Direction. | `inverse=True` | **yes** — a remap that transposes extents moves the patch grid | **yes** | **yes** — when the case's direction is a signed axis permutation; no on an oblique one (it is resampled) | | `Permute` | Permute spatial axes. `dims` is a pipe-separated axis list. | `dims="1\|0\|2", inverse=True` | **yes** | **yes** | **yes** — index remap | | `Flip` | Flip spatial axes. | `dims="1\|0\|2", inverse=True` | no | **yes** (self-inverse) | **yes** — index remap | diff --git a/docs/source/troubleshooting.md b/docs/source/troubleshooting.md index 2ec306b3..468bfe8b 100644 --- a/docs/source/troubleshooting.md +++ b/docs/source/troubleshooting.md @@ -103,7 +103,7 @@ depends on its arguments: pointwise covers `TensorCast`, `Clip` with **fixed** b and `Standardize` given **both** `mean` and `std`; `Normalize` and an automatic `Standardize` are global-statistic (still streamable, one statistics pass first); and the region kinds are `Flip`, `Permute`, **axis-aligned** `Canonical`, -`ResampleToShape`/`ResampleToResolution`, `Dilate` and `Gradient`. A `Clip` with a +`Resample` (to a grid, or through a stored transform or field), `Dilate` and `Gradient`. A `Clip` with a percentile bound needs the whole histogram, and `Padding` streams on the **write** side only — on the read side it loads the case whole. See the transform reference for the per-transform answer. diff --git a/docs/source/usage/large-images.md b/docs/source/usage/large-images.md index 7f93a62a..a6a5d0ac 100644 --- a/docs/source/usage/large-images.md +++ b/docs/source/usage/large-images.md @@ -83,8 +83,8 @@ requested output patch back to the source region on disk. | Halo | `Gradient`, `Dilate`, `Translate` | Enlarge the read and crop the result. | | Orientation | `Flip`, `Permute`, axis-aligned `Canonical` | Remap indices to the source. | | Crop | `Crop` once its box is known | Translate the target region. | -| Rescale | `ResampleToShape`, `ResampleToResolution` | Map through scale and add interpolation context. | -| Whole volume | masked transforms, histogram matching, arbitrary displacement | Use the bounded buffer. | +| Rescale | `Resample` — to a grid, or through a stored transform or field | Map through the grid change and add interpolation context. | +| Whole volume | masked transforms, histogram matching | Use the bounded buffer. | Region stages compose — any number of them, each pulling its read through the one before it — so a chain of remaps/halos/rescales still streams. An undeclared @@ -174,7 +174,7 @@ automatically — there is no flag to set. Each output slab is finalized and written to disk as soon as its patches complete, so peak RAM is one patch window instead of the whole volume. Geometry inverses stream too, and they **compose**: a `Canonical`/`Flip`/`Permute` inverse remaps each slab to its written region, a -`Padding` inverse crops it in flight, a `ResampleToResolution`/`ResampleToShape` +`Padding` inverse crops it in flight, a `spacing`/`shape` `Resample` inverse resamples back through a sliding window — chained in any number, each pulling through the next — so a huge output at ORIGINAL resolution is written slab by slab without ever being held whole. A masked finalize (`Mask`) streams diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 09f813d7..ad7a2655 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -1300,6 +1300,11 @@ def _target_from( "A resample writes on one grid: give its density (spacing), its extent (shape) or the" " image whose grid to adopt (reference) -- and only one of them.", ) + if reference and not str(reference).strip(): + raise TransformError( + "'Resample' was given a blank reference.", + "Name the entry whose grid to adopt: Resample: {reference: 822174, reference_group: Volume}.", + ) if align not in ("extent", "origin"): raise TransformError( f"'Resample' has an unknown align '{align}'.", @@ -1936,69 +1941,6 @@ def _target_is_own(self) -> bool: return isinstance(self._target, _OwnGrid) -class ResampleToResolution(Resample): - """Deprecated spelling of ``Resample: {spacing: ...}``.""" - - def __init__(self, spacing: list[float] = [1.0, 1.0, 1.0], inverse: bool = True) -> None: - super().__init__(spacing=spacing, inverse=inverse) - - -class ResampleToShape(Resample): - """Deprecated spelling of ``Resample: {shape: ...}``.""" - - def __init__(self, shape: list[int] = [100, 256, 256], inverse: bool = True) -> None: - super().__init__(shape=shape, inverse=inverse) - - -class ResampleToReference(Resample): - """Deprecated spelling of ``Resample: {reference: ...}``.""" - - def __init__( - self, - entry: str, - group: str | None = None, - dataset: str | None = None, - field: str | None = None, - field_group: str | None = None, - fill: float = 0.0, - interpolation: str | None = None, - inverse: bool = True, - ) -> None: - 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. Resample: {reference: 822174, reference_group: Volume}.", - ) - super().__init__( - reference=entry, - reference_group=group, - reference_dataset=dataset, - field=field, - field_group=field_group, - fill=fill, - interpolation=interpolation, - inverse=inverse, - ) - - -class ResampleTransform(Resample): - """Deprecated spelling of ``Resample: {transforms: ...}``.""" - - def __init__( - self, - transforms: dict[str, bool], - interpolation: str | None = None, - fill: float = 0.0, - inverse: bool = False, - ) -> None: - if not transforms: - raise TransformError( - "'ResampleTransform' needs at least one group of stored transforms to apply.", - "Name it and say whether to invert it, e.g. Resample: {transforms: {reg: false}}.", - ) - super().__init__(transforms=transforms, interpolation=interpolation, fill=fill, inverse=inverse) - - class Mask(Transform): """Set everything outside a mask to a constant. @@ -2479,18 +2421,6 @@ def check_bound(self, field: torch.Tensor, name: str) -> None: ) -class Warp(Resample): - """Deprecated spelling of ``Resample: {field: ...}`` — a warp on the case's own grid.""" - - def __init__(self, field: str, group: str | None = None, interpolation: str = "linear") -> None: - if not field or not str(field).strip(): - raise TransformError( - "'Warp' needs a 'field': the displacement field to resample through.", - "Declare it, e.g. Resample: {field: ./DVF:omezarr}.", - ) - super().__init__(field=field, field_group=group, interpolation=interpolation, inverse=False) - - class Reduce(Transform): """Fold every case of a group into one volume, at fixed voxel. @@ -2577,7 +2507,7 @@ class Expand(Transform): Clip: {min_value: 0.0, max_value: 400.0} # once per case Expand: {nb: 8, pattern: "{name}_r{a:02d}"} Rotate: {a_min: -15, a_max: 15} # a draw, per copy - ResampleToResolution: {spacing: [2, 2, 2]} # a transform, per copy + Resample: {spacing: [2, 2, 2]} # a transform, per copy Brightness: {b_std: 0.2} # another draw, per copy Write: {dataset: ./Augmented:omezarr} diff --git a/konfai/transformer.py b/konfai/transformer.py index eabd3a56..6daacf9c 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -767,7 +767,7 @@ def run_process(self, world_size: int, global_rank: int, local_rank: int, datalo shard = self._shards[global_rank] counts = {"STREAM": 0, "LOAD": 0, "WHOLE-VOLUME": 0, "SKIP": 0, "REDUCE": 0} # 'error' holds at run time too: a fallback the plan could not see (a sweep that fails, a - # Warp bound exceeded) raises at that case instead of quietly costing a volume. + # field bound exceeded) raises at that case instead of quietly costing a volume. allow_fallback = self.on_fallback != "error" def description() -> str: diff --git a/tests/integration/test_konfai_streamed_prediction.py b/tests/integration/test_konfai_streamed_prediction.py index 39b0ce31..f10960c8 100644 --- a/tests/integration/test_konfai_streamed_prediction.py +++ b/tests/integration/test_konfai_streamed_prediction.py @@ -21,7 +21,7 @@ The geometry variants exercise the write dispatcher end to end, one per region kind and then in composition: a ``Canonical`` inverse (ORIENTATION — in-slab mirrors), a ``Padding`` inverse (CROP), a -``ResampleToResolution`` inverse on a uint8 chain (REGRID, streamed in nearest mode, byte-exact) and on +a spacing ``Resample`` inverse on a uint8 chain (REGRID, streamed in nearest mode, byte-exact) and on a float chain (REGRID, streamed in linear mode, matching the reference to float-rounding), a two-inverse pipe, and the full three-inverse stack (crop + rescale + reorient composed, streamed end to end). The TTA variants exercise the slab-synchronized cross-copy reduce: an in-plane flip @@ -119,11 +119,11 @@ def main() -> None: inverse: true""", # REGRID: the inverse resamples back to the stored grid. "ResampleLabel": """ transforms: - ResampleToResolution: + Resample: spacing: [0.5, 0.5, -1.0] inverse: true""", "ResampleFloat": """ transforms: - ResampleToResolution: + Resample: spacing: [0.5, 0.5, -1.0] inverse: true""", # Several region stages compose into one streamed pipe: crop, then flip, straight to the sink. @@ -140,7 +140,7 @@ def main() -> None: "GeometryStack": """ transforms: Canonical: inverse: true - ResampleToResolution: + Resample: spacing: [0.5, 0.5, -1.0] inverse: true Padding: diff --git a/tests/integration/test_transform_example.py b/tests/integration/test_transform_example.py index 4fc61093..8784086b 100644 --- a/tests/integration/test_transform_example.py +++ b/tests/integration/test_transform_example.py @@ -56,7 +56,7 @@ def cohort(tmp_path_factory: pytest.TempPathFactory) -> Path: @pytest.mark.integration def test_the_template_example_folds_the_cohort_into_one_entry(cohort: Path) -> None: - """N to 1. The README's headline claim, and the reason ResampleToReference is in the chain: + """N to 1. The README's headline claim, and the reason a reference Resample is in the chain: the cases do not share a grid, so ``grid: strict`` would refuse them as stored.""" planned = _run([*konfai_cli_command(), "TRANSFORM", "--config", "Transform.yml", "--plan"], cohort) assert planned.returncode == 0, f"the template example does not plan:\n{planned.stdout}{planned.stderr}" diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py index 3842d5c4..d9fa5433 100644 --- a/tests/unit/test_api.py +++ b/tests/unit/test_api.py @@ -30,7 +30,7 @@ from konfai import api # noqa: E402 from konfai.data.reduction import Std # noqa: E402 -from konfai.data.transform import Clip, Magnitude, Resample, Warp, Write # noqa: E402 +from konfai.data.transform import Clip, Magnitude, Resample, Write # noqa: E402 from konfai.metric.measure import Dice # noqa: E402 from konfai.utils.errors import ConfigError, KonfAIError # noqa: E402 @@ -64,10 +64,15 @@ def test_a_repeated_mapping_stage_is_qualified_by_resolution() -> None: def test_a_subclass_delegating_to_super_keeps_its_own_spelling() -> None: - """``Warp(field=...)`` expands into ``Resample`` arguments internally; the recorded spelling is - the caller's, so the tree references ``Warp`` with the caller's kwargs and rebinds identically.""" - stage = Warp(field="./DVF:omezarr", group="DVF") - assert stage._konfai_given == {"field": "./DVF:omezarr", "group": "DVF"} + """The recorded spelling is the caller's: a subclass expanding into ``Resample`` arguments + inside ``super().__init__`` records its OWN kwargs, so the tree references the subclass with + what the caller wrote and rebinds identically.""" + + class FieldBeside(Resample): + def __init__(self, group: str) -> None: + super().__init__(field_group=group) + + assert FieldBeside(group="DVF")._konfai_given == {"group": "DVF"} def test_the_chain_tree_is_the_yaml_subtree() -> None: diff --git a/tests/unit/test_case_expansion.py b/tests/unit/test_case_expansion.py index 811fe9cd..3354889a 100644 --- a/tests/unit/test_case_expansion.py +++ b/tests/unit/test_case_expansion.py @@ -30,7 +30,7 @@ import torch from konfai.data.augmentation import Brightness, Flip, Permute, Scale from konfai.data.patching import DatasetManager -from konfai.data.transform import Clip, Expand, Save, TensorCast, Transform, Write, split_expand +from konfai.data.transform import Clip, Expand, Resample, Save, TensorCast, Transform, Write, split_expand from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import PatchError, TransformError @@ -467,7 +467,6 @@ def test_interleaved_patch_reads_of_two_copies_each_keep_their_own_grid(tmp_path a re-read of copy 1 after copy 2's plan silently returned copy 2's sampling of copy 1's data. """ from konfai.data.patching import DatasetPatch - from konfai.data.transform import Resample source = _source(tmp_path) permute = _draw(Permute(prob_permute=[0.5, 0.5])) diff --git a/tests/unit/test_itk_transforms.py b/tests/unit/test_itk_transforms.py index 0b6e0ef2..259ef182 100644 --- a/tests/unit/test_itk_transforms.py +++ b/tests/unit/test_itk_transforms.py @@ -56,11 +56,11 @@ def test_apply_to_data_transform_returns_ndarray() -> None: def test_resample_transform_applies_displacement_in_physical_space() -> None: - # ResampleTransform must not add the physical (dx, dy, dz) displacement straight onto a (z, y, x) + # A stored-transform Resample must not add the physical (dx, dy, dz) displacement straight onto a (z, y, x) # voxel-index grid: that transposes x/z and treats millimetres as voxels. A +6 mm translation along # X on a 2 mm-X grid must move content 3 voxels along X (not 6 voxels along Z). import torch - from konfai.data.transform import ResampleTransform + from konfai.data.transform import Resample from konfai.utils.dataset import Attribute volume = torch.zeros(1, 8, 8, 8, dtype=torch.uint8) @@ -79,7 +79,7 @@ def is_dataset_exist(self, group: str, name: str) -> bool: def read_transform(self, group: str, name: str) -> "sitk.Transform": return translation - transform = ResampleTransform({"reg": False}) + transform = Resample(transforms={"reg": False}) transform.datasets = [_TransformStore()] out = transform("case", volume, attribute) diff --git a/tests/unit/test_resample.py b/tests/unit/test_resample.py index 9f1225a0..859e4663 100644 --- a/tests/unit/test_resample.py +++ b/tests/unit/test_resample.py @@ -16,10 +16,8 @@ """One ``Resample``: which grid to write on, what map to write it through, and what that fixed. -``ResampleToResolution``, ``ResampleToShape``, ``ResampleToReference``, ``ResampleTransform`` and -``Warp`` were five stages answering two questions between them, each with its own sampler and its -own idea of where a voxel is. They are now five spellings of this one, and the tests here are for -what only became checkable once there was a single answer to check. +One stage, one sampler, one idea of where a voxel is -- the tests here are for what is only +checkable because there is a single answer to check. """ from __future__ import annotations @@ -29,14 +27,8 @@ import torch from konfai.data.transform import ( Resample, - ResampleToReference, - ResampleToResolution, - ResampleToShape, - ResampleTransform, - Warp, ) from konfai.utils.dataset import Attribute -from konfai.utils.errors import TransformError sitk = pytest.importorskip("SimpleITK") @@ -67,304 +59,6 @@ def _as_image(volume: np.ndarray, attribute: Attribute) -> sitk.Image: return image -# ------------------------------------------------------------------ one class, five spellings - - -@pytest.mark.parametrize( - ("alias", "unified"), - [ - (lambda: ResampleToResolution(spacing=[2.0, 1.5, 1.5]), lambda: Resample(spacing=[2.0, 1.5, 1.5])), - (lambda: ResampleToShape(shape=[12, 20, 22]), lambda: Resample(shape=[12, 20, 22])), - (lambda: ResampleToShape(shape=[0, 20, 0]), lambda: Resample(shape=[0, 20, 0])), - ], -) -def test_a_spelling_and_the_unified_stage_are_the_same_stage(alias, unified) -> None: - """The published names are argument translations, not behaviour of their own. - - Kept because they appear in shipped configs and in every bundle on the hub -- and kept THIN, - because a spelling that carries logic is a second implementation waiting to drift. - """ - volume = torch.from_numpy(_volume()) - left, right = alias(), unified() - assert isinstance(left, Resample) - assert left.apply_inverse == right.apply_inverse - - got = left("case", volume.clone(), _attributes()) - want = right("case", volume.clone(), _attributes()) - torch.testing.assert_close(got, want, rtol=0, atol=0) - - -def test_every_spelling_is_the_one_class() -> None: - for stage in ( - ResampleToResolution(), - ResampleToShape(), - ResampleToReference(entry="x"), - ResampleTransform(transforms={"reg": False}), - Warp(field="./x:h5", group="DVF"), - ): - assert isinstance(stage, Resample) - - -def test_the_three_ways_to_name_a_target_grid_are_exclusive() -> None: - with pytest.raises(TransformError, match="three ways to say the same thing"): - Resample(spacing=[1.0, 1.0, 1.0], shape=[4, 4, 4]) - - -# ------------------------------------------------------------------ where the new grid sits - - -def test_extent_alignment_keeps_the_field_of_view_and_origin_alignment_keeps_voxel_zero() -> None: - """The one silent choice in the family, made explicit. - - ``extent`` makes the outer faces coincide, which is ``F.interpolate``'s map and what KonfAI has - always done; ``origin`` keeps voxel zero's centre where it is, which is what resampling onto a - grid that shares an origin does. A quarter of a voxel of anatomy separates them. - """ - attribute = _attributes() - volume = torch.from_numpy(_volume()) - - Resample(spacing=[2.0, 1.5, 1.5], align="extent")("case", volume.clone(), attribute) - extent_origin = attribute.get_np_array("Origin").copy() - extent_spacing = attribute.get_np_array("Spacing").copy() - - attribute = _attributes() - Resample(spacing=[2.0, 1.5, 1.5], align="origin")("case", volume.clone(), attribute) - - np.testing.assert_allclose(attribute.get_np_array("Origin"), _ORIGIN) - np.testing.assert_allclose(attribute.get_np_array("Spacing"), [2.0, 1.5, 1.5]) - # Extent alignment puts voxel zero half the spacing change away, on every axis. - np.testing.assert_allclose(extent_origin, np.asarray(_ORIGIN) + 0.5 * (extent_spacing - np.asarray(_SPACING))) - - -def test_an_unknown_alignment_is_refused_at_construction() -> None: - with pytest.raises(TransformError, match="unknown align"): - Resample(spacing=[1.0, 1.0, 1.0], align="corners") - - -# ------------------------------------------------------------------ the header describes the data - - -@pytest.mark.parametrize( - "kwargs", - [ - {"spacing": [2.0, 1.5, 1.5]}, - {"spacing": [2.0, 1.5, 1.5], "align": "origin"}, - {"shape": [12, 20, 22]}, - {"spacing": [0.9, 1.7, 1.1]}, - ], -) -def test_the_recorded_header_describes_the_grid_that_was_actually_sampled(kwargs) -> None: - """The check neither predecessor could pass, because neither wrote a placement at all. - - ``ResampleToResolution`` recorded the spacing that was ASKED FOR while sampling at ``n_in/n_out`` - times the source's -- up to a millimetre of drift across a volume -- and left the Origin alone - while sampling half a spacing-change away from it. Nothing downstream could see either: the - voxels are all real, and the header is the only witness. - - So the oracle is built FROM THE RECORDED HEADER. If the two disagree, resampling the source onto - the grid the header claims cannot reproduce the voxels the stage returned. - """ - volume = _volume() - attribute = _attributes() - got = Resample(**kwargs)("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] - - grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) - grid.SetOrigin(attribute.get_np_array("Origin").tolist()) - grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) - grid.SetDirection(attribute.get_np_array("Direction").tolist()) - want = sitk.GetArrayFromImage( - sitk.Resample(_as_image(volume, _attributes()), grid, sitk.Transform(), sitk.sitkLinear, 0.0) - ) - np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) - - -def test_an_oblique_case_records_an_oblique_header() -> None: - """A direction is carried through, not quietly dropped -- and the data follows it.""" - angle = np.deg2rad(23.0) - cos, sin = float(np.cos(angle)), float(np.sin(angle)) - direction = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]) - volume = _volume() - attribute = _attributes(direction=direction) - got = Resample(spacing=[2.0, 1.5, 1.5])("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] - - np.testing.assert_allclose(attribute.get_np_array("Direction").reshape(3, 3), direction) - grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) - grid.SetOrigin(attribute.get_np_array("Origin").tolist()) - grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) - grid.SetDirection(direction.reshape(-1).tolist()) - want = sitk.GetArrayFromImage( - sitk.Resample(_as_image(volume, _attributes(direction=direction)), grid, sitk.Transform(), sitk.sitkLinear, 0.0) - ) - np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) - - -# ------------------------------------------------------------------ an image and its label map - - -def test_a_label_map_lands_on_the_same_voxels_as_the_image_beside_it() -> None: - """The bug that had no symptom: a mask resampled with its CT came out shifted against it. - - ``F.interpolate``'s nearest reads ``floor(o * scale)`` where its linear reads - ``scale * (o + 0.5) - 0.5`` -- so the label map lagged the image of the SAME stage by - ``(scale - 1) / 2`` source voxels. At 0.5 mm resampled to 3 mm that is 2.5 source voxels, 1.25 mm - of anatomy, and both volumes are entirely plausible on their own. - - A ramp makes it visible: linear interpolation of a linear function is exact, so the resampled - image IS the continuous source coordinate, and the resampled label map must be its rounding. - """ - extent = 48 - ramp = np.broadcast_to(np.arange(extent, dtype=np.float32).reshape(1, 1, -1), (1, 4, 4, extent)) - attribute = _attributes(origin=[0.0, 0.0, 0.0], spacing=[0.5, 1.0, 1.0]) - - image = Resample(spacing=[3.0, 1.0, 1.0])("case", torch.from_numpy(np.ascontiguousarray(ramp)), attribute) - labels = Resample(spacing=[3.0, 1.0, 1.0])( - "case", torch.from_numpy(np.ascontiguousarray(ramp).astype(np.uint8)), _attributes([0.0] * 3, [0.5, 1.0, 1.0]) - ) - - coordinate = image.numpy()[0, 0, 0] - picked = labels.numpy()[0, 0, 0].astype(np.int64) - # The edges clamp, so the interior is where the ramp still reads its own coordinate. - interior = slice(1, -1) - np.testing.assert_array_equal(picked[interior], np.floor(coordinate[interior] + 0.5).astype(np.int64)) - assert float(np.abs(picked[interior] - coordinate[interior]).max()) <= 0.5 - - -# ------------------------------------------------------------------ the count - - -def test_a_spacing_that_binary_cannot_hold_does_not_lose_a_slice() -> None: - """90 voxels of 0.7 mm re-cut at 1.5 mm is 42.0 -- and in float64 it is 41.999999999999997. - - Truncating that gives 41: one slice of anatomy dropped, and a recorded spacing that no longer - covers what was read. The old float32 round-trip happened to land above; nothing said so. - """ - attribute = _attributes(origin=[0.0] * 3, spacing=[0.7, 0.7, 0.7]) - shape = Resample(spacing=[1.5, 1.5, 1.5]).transform_shape("", "case", [90, 90, 90], attribute) - assert shape == [42, 42, 42] - - -# ------------------------------------------------------------------ one grid change, one map - - -def test_a_change_of_grid_and_a_warp_are_one_interpolation(tmp_path) -> None: - """Asked for together they compose into one coordinate per voxel, checked against sitk's own. - - Two stages would interpolate the same voxels twice, and the second pass invents none of what the - first smoothed away -- which is the whole reason an atlas's appearance is rebuilt from native - volumes rather than from warped ones. - """ - from konfai.utils.dataset import Dataset - - field_shape, field_origin, field_spacing = (6, 8, 9), [10.0, -5.0, 38.0], [4.0, 3.5, 3.0] - field = np.zeros((3, *field_shape), dtype=np.float32) - for component, value in enumerate((1.5, -2.0, 0.75)): - field[component] = value - - store = Dataset(tmp_path / "DVF", "h5") - store.write("DVF", "case", field, _attributes(field_origin, field_spacing)) - - volume = _volume() - attribute = _attributes() - stage = Resample(spacing=[2.0, 1.5, 1.5], field=str(tmp_path / "DVF") + ":h5", field_group="DVF") - stage.set_datasets([store]) - got = stage("case", torch.from_numpy(volume.copy()), attribute).numpy()[0] - - grid = sitk.Image(*reversed(got.shape), sitk.sitkFloat32) - grid.SetOrigin(attribute.get_np_array("Origin").tolist()) - grid.SetSpacing(attribute.get_np_array("Spacing").tolist()) - vector = sitk.GetImageFromArray(np.moveaxis(field, 0, -1).astype(np.float64), isVector=True) - vector.SetOrigin(field_origin) - vector.SetSpacing(field_spacing) - want = sitk.GetArrayFromImage( - sitk.Resample( - _as_image(volume, _attributes()), - grid, - sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)), - sitk.sitkLinear, - 0.0, - ) - ) - np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-3) - - -# ------------------------------------------------------------------ which loop runs - - -def test_a_map_that_factorises_takes_the_separable_loop() -> None: - """The optimisation needs a test, or it can be lost to a refactor with everything still green. - - A grid change between axis-aligned volumes reads one axis at a time; a rotation between them, or - a displacement, cannot and falls to the coordinate volume. Both are correct — the difference is - 43x versus 660x of ``F.interpolate`` on a CT-sized case, which no assertion about values shows. - """ - from konfai.data.geometry import Grid - from konfai.data.sampling import separable_source_index - - device = torch.device("cpu") - source = Grid(_SHAPE, np.asarray(_ORIGIN), np.asarray(_SPACING), np.eye(3)) - aligned = source.resampled(spacing_xyz=np.asarray([2.0, 1.5, 1.5])) - assert separable_source_index(aligned, source, (), device) is not None - - angle = np.deg2rad(23.0) - cos, sin = float(np.cos(angle)), float(np.sin(angle)) - turned = Grid( - aligned.size_zyx, - aligned.origin_xyz, - aligned.spacing_xyz, - np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]), - ) - assert separable_source_index(turned, source, (), device) is None, "a rotation does not factorise" - - # A flip is still axis-aligned, so it factorises -- the test that the check is not merely - # "is the direction the identity". - flipped = Grid(source.size_zyx, source.origin_xyz, source.spacing_xyz, np.diag([-1.0, -1.0, 1.0])) - assert separable_source_index(aligned, flipped, (), device) is not None - - -def test_the_two_loops_agree_where_both_can_serve_the_same_map() -> None: - """Different summation orders, same answer to float rounding — and the same fill, exactly.""" - from konfai.data.geometry import Grid - from konfai.data.sampling import gather, gather_separable, separable_source_index, source_index - - device = torch.device("cpu") - volume = torch.from_numpy(_volume()) - source = Grid(_SHAPE, np.asarray(_ORIGIN), np.asarray(_SPACING), np.eye(3)) - # Placed so part of the target reaches past the case, which is where the fill rule shows. - target = Grid((14, 18, 20), np.asarray(_ORIGIN) - 4.0, np.asarray([2.0, 1.6, 1.6]), np.eye(3)) - - axes = separable_source_index(target, source, (), device) - assert axes is not None - fast = gather_separable(volume, axes, [0, 0, 0], list(_SHAPE), "linear", -999.0) - general = gather(volume, source_index(target, source, (), device), [0, 0, 0], list(_SHAPE), "linear", -999.0) - - np.testing.assert_array_equal((fast == -999.0).numpy(), (general == -999.0).numpy()) - torch.testing.assert_close(fast, general, rtol=1e-6, atol=1e-5) - - -def test_the_blend_order_is_the_same_for_a_region_as_for_the_whole_volume() -> None: - """Axes are blended most-reduced-first, and the key must not be the extents in hand. - - Blending an axis reduces it before the next one reads it, so the order decides how much data - every later pass moves -- 9x on a thick-slice CT brought to isotropic. But the order also decides - the SUMMATION order, so a region that chose differently from the whole volume would stop being - bit-identical to it, which is the one equality the streaming design rests on. Keyed on the two - grids' spacings, which a region shares with its volume, and not on their extents, which it does - not. - """ - from konfai.data.geometry import Grid - from konfai.data.sampling import blend_order - - source = Grid((64, 512, 512), np.zeros(3), np.asarray([0.7, 0.7, 3.0]), np.eye(3)) - target = source.resampled(spacing_xyz=np.asarray([1.0, 1.0, 1.0])) - whole = blend_order(target, source) - - # z triples while y and x shrink, so y and x are blended first. - assert whole == [1, 2, 0] - for start, stop in ((0, 8), (17, 41), (target.size_zyx[0] - 3, target.size_zyx[0])): - region = target.sub_grid((slice(start, stop), slice(0, target.size_zyx[1]), slice(0, target.size_zyx[2]))) - assert blend_order(region, source) == whole - - def test_an_axis_the_map_leaves_alone_is_left_alone() -> None: """A resample of one axis reads the other two, it does not blend them — and says so in the values. diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 14f68050..8e3b6c61 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -35,7 +35,7 @@ from konfai.data.geometry import AffineMap, Grid, TransformBound from konfai.data.patching import DatasetManager, DatasetPatch from konfai.data.sampling import source_window -from konfai.data.transform import LocalityKind, Reduce, ResampleToReference, ResampleToShape, Write +from konfai.data.transform import LocalityKind, Reduce, Resample, Write from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ConfigError, TransformError from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE @@ -79,14 +79,14 @@ def dataset(tmp_path: Path) -> Dataset: 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] +def _stage(dataset: Dataset, **kwargs: object) -> Resample: + arguments: dict[str, object] = {"reference": _CASE, "reference_group": "Reference", "fill": _FILL, **kwargs} + stage = Resample(**arguments) # type: ignore[arg-type] stage.set_datasets([dataset]) return stage -def _manager(dataset: Dataset, stage: ResampleToReference, group: str = "Case") -> DatasetManager: +def _manager(dataset: Dataset, stage: Resample, group: str = "Case") -> DatasetManager: return DatasetManager( index=0, group_src=group, @@ -273,7 +273,7 @@ def managers(with_stage: bool) -> list[DatasetManager]: for index in range(len(shapes)): stages: list[object] = [] if with_stage: - stage = ResampleToReference(entry="GRID", group="Reference", fill=_FILL) + stage = Resample(reference="GRID", reference_group="Reference", fill=_FILL) stage.set_datasets([dataset]) stages.append(stage) built.append( @@ -388,7 +388,7 @@ def test_it_is_refused_as_a_patch_transform(dataset: Dataset, monkeypatch: pytes # A change of extent needs no geometry at all, so it reaches config time as what it is. with pytest.raises(ConfigError, match="onto another grid"): - _check_patch_transform_locality(ResampleToShape(shape=[4, 4, 4]), "CT", "CT") + _check_patch_transform_locality(Resample(shape=[4, 4, 4]), "CT", "CT") def test_a_differing_direction_is_resampled_and_not_refused(tmp_path: Path) -> None: @@ -439,23 +439,23 @@ def test_a_case_that_never_meets_the_reference_is_refused(tmp_path: Path) -> Non def test_an_unknown_entry_is_refused(dataset: Dataset) -> None: - stage = ResampleToReference(entry="NOT_THERE", group="Reference") + stage = Resample(reference="NOT_THERE", reference_group="Reference") stage.set_datasets([dataset]) with pytest.raises(TransformError, match="cannot find reference 'NOT_THERE'"): stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) 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) + """Guessing which group holds the reference is not a guess worth making.""" + stage = Resample(reference=_CASE) stage.set_datasets([dataset]) with pytest.raises(TransformError, match="cannot tell which group"): stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) -def test_an_empty_entry_is_refused_at_construction() -> None: - with pytest.raises(TransformError, match="needs an 'entry'"): - ResampleToReference(entry=" ") +def test_a_blank_reference_is_refused_at_construction() -> None: + with pytest.raises(TransformError, match="blank reference"): + Resample(reference=" ", reference_group="Reference") # --------------------------------------------------------------------- what it announces @@ -482,7 +482,7 @@ def test_a_case_that_fills_the_grid_says_nothing(tmp_path: Path) -> None: 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 = Resample(reference=_CASE, reference_group="Reference") inside.set_datasets([dataset]) assert inside.plan_note("Case_out", _CASE, [20, 20, 20], source) is None @@ -577,16 +577,16 @@ def warped(tmp_path: Path) -> tuple[Dataset, Dataset, np.ndarray]: return images, fields, volume -def _warping(images: Dataset, fields: Dataset, **kwargs: object) -> ResampleToReference: +def _warping(images: Dataset, fields: Dataset, **kwargs: object) -> Resample: arguments: dict[str, object] = { - "entry": _CASE, - "group": "Reference", + "reference": _CASE, + "reference_group": "Reference", "field": f"{fields.filename}:h5", "field_group": "DVF", "fill": _FILL, **kwargs, } - stage = ResampleToReference(**arguments) # type: ignore[arg-type] + stage = Resample(**arguments) # type: ignore[arg-type] stage.set_datasets([images]) return stage @@ -734,9 +734,9 @@ def test_the_field_components_are_not_reversed(tmp_path: Path) -> None: uniform[0] = 2.0 fields.write("DVF", _CASE, uniform, geometry) - stage = ResampleToReference( - entry=_CASE, - group="Reference", + stage = Resample( + reference=_CASE, + reference_group="Reference", field=f"{fields.filename}:h5", field_group="DVF", fill=0.0, @@ -871,8 +871,8 @@ def test_a_field_beyond_its_recorded_bound_is_refused( attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.asarray([0.5, 0.5, 0.5]) lying = Dataset(tmp_path / "lying", "mha") lying.write("DVF", _CASE, _displacement(), attributes) - stage = ResampleToReference( - entry=_CASE, group="Reference", field=f"{tmp_path / 'lying'}:mha", field_group="DVF", fill=_FILL + stage = Resample( + reference=_CASE, reference_group="Reference", field=f"{tmp_path / 'lying'}:mha", field_group="DVF", fill=_FILL ) stage.set_datasets([images]) with pytest.raises(TransformError, match="displaces up to"): @@ -880,7 +880,7 @@ def test_a_field_beyond_its_recorded_bound_is_refused( def test_a_field_with_no_bound_still_streams(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: - """Warp's rule, and for the same reason: the run sizes each region's pull from the field + """The run sizes each region's pull from the field values it reads for sampling, so a missing recorded bound is a pricing gap, not a fallback.""" images, fields, _volume = warped stage = _warping(images, fields) @@ -891,7 +891,7 @@ def test_a_field_with_no_bound_still_streams(warped: tuple[Dataset, Dataset, np. def _stage_regrid_kind(images: Dataset) -> LocalityKind: - stage = ResampleToReference(entry=_CASE, group="Reference") + stage = Resample(reference=_CASE, reference_group="Reference") stage.set_datasets([images]) return stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind @@ -928,7 +928,7 @@ def test_fields_can_live_beside_the_cases(warped: tuple[Dataset, Dataset, np.nda """ images, fields, volume = warped images.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) - beside = ResampleToReference(entry=_CASE, group="Reference", field_group="DVF", fill=_FILL) + beside = Resample(reference=_CASE, reference_group="Reference", field_group="DVF", fill=_FILL) beside.set_datasets([images]) got = beside(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) @@ -952,7 +952,7 @@ def test_the_recorded_bound_is_read_from_every_root_not_the_first(tmp_path: Path 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") + stage = Resample(reference="CASE_1", reference_group="Reference", field_group="DVF") stage.set_datasets([first, second]) # 9.0 from the second root, not 1.0 from the first. diff --git a/tests/unit/test_resample_transform.py b/tests/unit/test_resample_transform.py index 879fa868..0394cccb 100644 --- a/tests/unit/test_resample_transform.py +++ b/tests/unit/test_resample_transform.py @@ -14,7 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""``ResampleTransform`` streaming: against SimpleITK, and against its own whole-volume path. +"""``Resample`` through stored transforms: against SimpleITK, and against its own whole-volume path. The fixture is high-frequency and its direction cosines are oblique, because a smooth phantom on an axis-aligned grid passes a map that is wrong in exactly the ways this stage can be wrong. @@ -23,7 +23,7 @@ import numpy as np import pytest import torch -from konfai.data.transform import LocalityKind, RegionContext, ResampleTransform +from konfai.data.transform import LocalityKind, RegionContext, Resample from konfai.utils.dataset import Attribute from konfai.utils.errors import TransformError @@ -59,7 +59,7 @@ def _attribute(image: "sitk.Image") -> Attribute: class _StoredTransform: - """The smallest thing that answers what ``ResampleTransform`` asks of a ``Dataset``.""" + """The smallest thing that answers what a stored-transform ``Resample`` asks of a ``Dataset``.""" def __init__(self, group: str, transform: "sitk.Transform") -> None: self.group = group @@ -112,8 +112,8 @@ def _families(image) -> list[tuple[str, "sitk.Transform"]]: ] -def _stage(image, transform, **kwargs) -> ResampleTransform: - stage = ResampleTransform(transforms={"reg": False}, **kwargs) +def _stage(image, transform, **kwargs) -> Resample: + stage = Resample(transforms={"reg": False}, **kwargs) stage.set_datasets([_StoredTransform("reg", transform)]) stage.transform_shape("", CASE, list(SIZE), _attribute(image)) return stage @@ -168,7 +168,7 @@ def test_the_streamed_slabs_agree_with_the_whole_volume(device: torch.device, ro torch.testing.assert_close(streamed, reference, rtol=0.0, atol=1e-5 * span, msg=label) -def _slab_reads(stage: ResampleTransform, attribute: Attribute, rows: int) -> list[int]: +def _slab_reads(stage: Resample, attribute: Attribute, rows: int) -> list[int]: reads = [] for start in range(0, SIZE[0], rows): target = (slice(start, min(start + rows, SIZE[0])), slice(0, SIZE[1]), slice(0, SIZE[2])) @@ -213,7 +213,7 @@ def test_a_boundable_cohort_declares_regrid(self): def test_a_case_without_geometry_falls_back_and_says_why(self): image = _image() - stage = ResampleTransform(transforms={"reg": False}) + stage = Resample(transforms={"reg": False}) stage.set_datasets([_StoredTransform("reg", _euler(image))]) stage.transform_shape("", CASE, list(SIZE), Attribute()) # no Origin/Spacing/Direction locality = stage.patch_locality(Attribute()) # judged on the header handed over @@ -227,7 +227,7 @@ def test_a_missing_transform_refuses_at_plan_time_and_says_which_group(self): the whole-volume path needs the same decode this refusal comes from. """ image = _image() - stage = ResampleTransform(transforms={"absent": False}) + stage = Resample(transforms={"absent": False}) stage.set_datasets([_StoredTransform("reg", _euler(image))]) with pytest.raises(TransformError, match="absent"): stage.transform_shape("", CASE, list(SIZE), _attribute(image)) @@ -245,14 +245,14 @@ def test_a_spline_order_with_no_kernel_refuses_at_plan_time(self): size = np.asarray(quadratic.GetParameters()).size quadratic.SetParameters(list(np.random.RandomState(3).uniform(-5.0, 5.0, size))) - stage = ResampleTransform(transforms={"reg": False}) + stage = Resample(transforms={"reg": False}) stage.set_datasets([_StoredTransform("reg", quadratic)]) with pytest.raises(TransformError, match="order 2"): stage.transform_shape("", CASE, list(SIZE), _attribute(image)) def test_inverting_a_spline_refuses_at_plan_time_with_the_remedy(self): image = _image() - stage = ResampleTransform(transforms={"reg": True}) + stage = Resample(transforms={"reg": True}) stage.set_datasets([_StoredTransform("reg", _bspline(image))]) with pytest.raises(TransformError, match="Store the inverse"): stage.transform_shape("", CASE, list(SIZE), _attribute(image)) @@ -260,7 +260,7 @@ def test_inverting_a_spline_refuses_at_plan_time_with_the_remedy(self): def test_inverting_a_rigid_map_is_exact_and_still_streams(self): image = _image() transform = _euler(image) - stage = ResampleTransform(transforms={"reg": True}) + stage = Resample(transforms={"reg": True}) stage.set_datasets([_StoredTransform("reg", transform)]) stage.transform_shape("", CASE, list(SIZE), _attribute(image)) assert stage.patch_locality(_attribute(image)).kind is LocalityKind.REGRID @@ -296,14 +296,14 @@ def test_the_fill_reaches_where_the_map_leaves_the_source(self): class TestRefusals: def test_an_unknown_interpolation_is_refused_at_construction(self): with pytest.raises(TransformError, match="interpolation"): - ResampleTransform(transforms={"reg": False}, interpolation="bspline") + Resample(transforms={"reg": False}, interpolation="bspline") def test_no_transforms_is_refused_at_construction(self): - with pytest.raises(TransformError, match="at least one group"): - ResampleTransform(transforms={}) + with pytest.raises(TransformError, match="empty 'transforms'"): + Resample(transforms={}) def test_the_inverse_direction_says_what_to_do_instead(self): - stage = ResampleTransform(transforms={"reg": False}) + stage = Resample(transforms={"reg": False}) assert stage.inverse_patch_locality(Attribute()).kind is LocalityKind.WHOLE_VOLUME with pytest.raises(TransformError, match="inverse: false"): stage.inverse(CASE, torch.zeros(1, 2, 2, 2), Attribute()) @@ -317,7 +317,7 @@ def test_two_groups_compose_as_they_always_did(self): image = _image(oblique=False) first = sitk.TranslationTransform(3, (4.0, 0.0, 0.0)) second = sitk.ScaleTransform(3, (1.5, 1.5, 1.5)) - stage = ResampleTransform(transforms={"a": False, "b": False}) + stage = Resample(transforms={"a": False, "b": False}) class _Two: def is_dataset_exist(self, group: str, name: str) -> bool: diff --git a/tests/unit/test_streamed_read_dispatcher.py b/tests/unit/test_streamed_read_dispatcher.py index 3c7d46bc..c64799ee 100644 --- a/tests/unit/test_streamed_read_dispatcher.py +++ b/tests/unit/test_streamed_read_dispatcher.py @@ -48,7 +48,6 @@ Permute, RegionContext, Resample, - ResampleToShape, Softmax, TensorCast, Transform, @@ -162,9 +161,7 @@ def test_stream_composed_rescale_and_orientation_matches_whole_volume(assert_str # window, on the RESAMPLED grid the fold computed between them. rng = np.random.default_rng(7) volume = (rng.standard_normal((1, 8, 8)).astype(np.float32)) * 100.0 - manager = assert_stream_matches_whole_volume( - volume, [ResampleToShape(shape=[12, 12]), Flip("0")], [4, 4], atol=1e-3 - ) + manager = assert_stream_matches_whole_volume(volume, [Resample(shape=[12, 12]), Flip("0")], [4, 4], atol=1e-3) plans = manager._resolve_patch_stream_source(0, True).stage_plans assert [plan.kind.value for plan in plans] == ["regrid", "orientation"] assert tuple(plans[1].in_shape) == (12, 12) @@ -175,7 +172,7 @@ def test_stream_composed_triple_region_chain_matches_whole_volume(assert_stream_ rng = np.random.default_rng(11) volume = (rng.standard_normal((1, 8, 6)).astype(np.float32)) * 100.0 manager = assert_stream_matches_whole_volume( - volume, [Flip("0"), ResampleToShape(shape=[12, 9]), Permute("1|0")], [4, 4], atol=1e-3 + volume, [Flip("0"), Resample(shape=[12, 9]), Permute("1|0")], [4, 4], atol=1e-3 ) plans = manager._resolve_patch_stream_source(0, True).stage_plans assert [plan.kind.value for plan in plans] == ["orientation", "regrid", "orientation"] @@ -456,10 +453,10 @@ def test_stream_resample_nearest_strong_downsampling_matches_whole_volume(build_ volume = (np.arange(1 * 40 * 40).reshape(1, 40, 40) % 7).astype(np.uint8) shape = [6, 6] patch = [3, 3] - stream_manager = build_streaming_manager(volume, [ResampleToShape(shape=shape)], patch) + stream_manager = build_streaming_manager(volume, [Resample(shape=shape)], patch) assert stream_manager.can_stream_patch(0) - reference_manager = build_streaming_manager(volume, [ResampleToShape(shape=shape)], patch) + reference_manager = build_streaming_manager(volume, [Resample(shape=shape)], patch) reference_manager.load(reference_manager.transforms, [], load_augmentations=False) size = stream_manager.patch.get_size(0) @@ -474,7 +471,7 @@ def test_stream_resample_nearest_strong_downsampling_matches_whole_volume(build_ def test_streamed_nearest_resample_matches_whole_volume_at_any_ratio(n_in: int, n_out: int) -> None: """The streamed nearest gather must pick the same source voxel as F.interpolate, per axis.""" volume = (torch.arange(n_in**3, dtype=torch.int32) % 251).to(torch.uint8).reshape(1, n_in, n_in, n_in) - resample = ResampleToShape(shape=[n_out, n_out, n_out], inverse=False) + resample = Resample(shape=[n_out, n_out, n_out], inverse=False) attribute = Attribute() attribute["Spacing"] = np.ones(3) expected = resample("case", volume.clone(), Attribute(attribute)) @@ -490,7 +487,7 @@ def test_streamed_nearest_resample_matches_whole_volume_at_any_ratio(n_in: int, def test_streamed_resample_handles_2d(dtype: torch.dtype) -> None: """The gather must not assume three spatial axes.""" volume = (torch.arange(1 * 9 * 11, dtype=torch.float32).reshape(1, 9, 11) % 17).to(dtype) - resample = ResampleToShape(shape=[5, 6], inverse=False) + resample = Resample(shape=[5, 6], inverse=False) attribute = Attribute() attribute["Spacing"] = np.ones(2) expected = resample("case", volume.clone(), Attribute(attribute)) @@ -539,7 +536,7 @@ def test_replanning_after_epoch_redraw_keeps_the_stored_geometry(build_streaming reorienting from the second epoch on, while the geometry keys stack once more per epoch. """ volume = np.arange(1 * 8 * 8 * 8, dtype=np.float32).reshape(1, 8, 8, 8) - transform = ResampleToShape(shape=[16, 16, 16]) if transform_case == "resample" else Canonical() + transform = Resample(shape=[16, 16, 16]) if transform_case == "resample" else Canonical() streamed = build_streaming_manager(volume, [transform], [4, 4, 4], _flip_augmentations()) assert streamed.can_stream_patch(0) diff --git a/tests/unit/test_streamed_write_dispatcher.py b/tests/unit/test_streamed_write_dispatcher.py index e93a01b5..affe9a7d 100644 --- a/tests/unit/test_streamed_write_dispatcher.py +++ b/tests/unit/test_streamed_write_dispatcher.py @@ -42,7 +42,7 @@ Padding, Permute, RegionContext, - ResampleToResolution, + Resample, Softmax, Standardize, TensorCast, @@ -220,7 +220,7 @@ def test_stream_rescale_nearest_is_byte_identical_to_the_whole_volume_inverse(se # to be smaller: byte-identical by construction rather than by agreement. rng = np.random.default_rng(seed) volume = torch.from_numpy(rng.integers(0, 7, size=(C, Z, Y, X)).astype(np.uint8)) - resample = ResampleToResolution([1.0, 1.0, 1.0]) + resample = Resample([1.0, 1.0, 1.0]) attribute = Attribute() attribute["Spacing"] = torch.tensor([1.0, 1.0, 1.0]) attribute["Size"] = np.asarray([12, 9, 7]) # the size the inverse restores @@ -251,7 +251,7 @@ def test_stream_rescale_linear_matches_the_whole_volume_inverse_to_float_roundin # sample in the same place. There is no tolerance to negotiate here any more. rng = np.random.default_rng(seed) volume = torch.from_numpy(rng.standard_normal((C, Z, Y, X)).astype(np.float32)) * 100.0 - resample = ResampleToResolution([1.0, 1.0, 1.0]) + resample = Resample([1.0, 1.0, 1.0]) attribute = Attribute() attribute["Spacing"] = torch.tensor([1.0, 1.0, 1.0]) attribute["Size"] = np.asarray([12, 9, 7]) @@ -306,7 +306,7 @@ def test_stream_window_is_bounded_by_the_pull_span() -> None: # The whole point: the buffer holds the pull span of the pending output rows, never the volume. tall = 64 volume = (torch.arange(1 * tall * Y * X).reshape(1, tall, Y, X) % 5).to(torch.uint8) - resample = ResampleToResolution([1.0, 1.0, 1.0]) + resample = Resample([1.0, 1.0, 1.0]) attribute = Attribute() attribute["Spacing"] = torch.tensor([1.0, 1.0, 1.0]) attribute["Size"] = np.asarray([96, Y, X]) @@ -364,11 +364,11 @@ def test_inverse_locality_defaults_and_overrides() -> None: # Geometry inverses declare their own kind. assert Padding().inverse_patch_locality(empty).kind is LocalityKind.CROP # A resample inverse is patch-native only when the Size stack it pops is on the case. - assert ResampleToResolution().inverse_patch_locality(empty).kind is LocalityKind.WHOLE_VOLUME + assert Resample().inverse_patch_locality(empty).kind is LocalityKind.WHOLE_VOLUME seeded = Attribute() seeded["Size"] = np.asarray([4, 4, 4]) seeded["Size"] = np.asarray([2, 2, 2]) - assert ResampleToResolution().inverse_patch_locality(seeded).kind is LocalityKind.REGRID + assert Resample().inverse_patch_locality(seeded).kind is LocalityKind.REGRID # Canonical judges the POPPED state: without a stacked direction there is nothing to invert onto. assert Canonical().inverse_patch_locality(empty).kind is LocalityKind.WHOLE_VOLUME @@ -613,14 +613,14 @@ def test_add_layer_streams_a_forward_region_final_transform(tmp_path, monkeypatc def test_add_layer_streams_a_full_geometry_stack_through_the_composed_pipe(tmp_path, monkeypatch) -> None: - # The general case the composition exists for: Canonical + ResampleToResolution + Padding forward, + # The general case the composition exists for: Canonical + Resample + Padding forward, # so the finalize chain carries CROP + REGRID + ORIENTATION in sequence. With the labelmap cast # to uint8 before the reduction, the whole stack streams to the sink and must match the # whole-volume path bit for bit. volume = (torch.arange(1 * 6 * 4 * 3).reshape(1, 6, 4, 3) % 5).to(torch.float32) transforms = [ Canonical(inverse=True), - ResampleToResolution([0.5, 0.5, -1.0], inverse=True), + Resample([0.5, 0.5, -1.0], inverse=True), Padding([0, 0, 0, 0, 2, 1], inverse=True), ] before = [TensorCast("uint8", inverse=False)] diff --git a/tests/unit/test_transform.py b/tests/unit/test_transform.py index 468c5689..9347e893 100644 --- a/tests/unit/test_transform.py +++ b/tests/unit/test_transform.py @@ -15,7 +15,7 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for ``konfai.data.transform``: Clip, Dilate, Norm, Crop, Standardize, Padding, -ResampleToResolution/ResampleToShape, InferenceStack, and KonfAIInference.""" +Resample, InferenceStack, and KonfAIInference.""" import os import sys @@ -41,8 +41,7 @@ OneHot, Padding, Reduce, - ResampleToResolution, - ResampleToShape, + Resample, Squeeze, StandardDeviation, Standardize, @@ -287,25 +286,25 @@ def test_padding_after_the_data_keeps_origin(image_attributes): # -------------------------------------------------------------------------------------- -# ResampleToResolution / ResampleToShape +# Resample: spacing / shape # -------------------------------------------------------------------------------------- def test_resample_to_resolution_transform_shape_missing_spacing_raises(): """A density change is meaningless without the density it starts from, so it refuses.""" with pytest.raises(TransformError): - ResampleToResolution().transform_shape("group", "case", [10, 10, 10], Attribute()) + Resample(spacing=[1.0, 1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], Attribute()) def test_resample_to_shape_needs_no_spacing_at_all(): """A count is a count. Only a DENSITY change needs the density it starts from. - This used to refuse alongside ``ResampleToResolution``, which cost nothing but told the user to + Refusing here would cost nothing but would tell the user to go and find a geometry for an operation that is a pure resize. With no header the grid is the identity, and the map degenerates to the size ratio it always was. """ - assert ResampleToShape(shape=[4, 5, 6]).transform_shape("group", "case", [10, 10, 10], Attribute()) == [4, 5, 6] - resized = ResampleToShape(shape=[4, 5, 6])("case", torch.zeros(1, 10, 10, 10), Attribute()) + assert Resample(shape=[4, 5, 6]).transform_shape("group", "case", [10, 10, 10], Attribute()) == [4, 5, 6] + resized = Resample(shape=[4, 5, 6])("case", torch.zeros(1, 10, 10, 10), Attribute()) assert list(resized.shape[1:]) == [4, 5, 6] @@ -314,17 +313,17 @@ def test_resample_to_resolution_transform_shape_dimension_mismatch_message(): attributes = Attribute() attributes["Spacing"] = np.asarray([1.0, 1.0], dtype=np.float64) with pytest.raises(TransformError) as excinfo: - ResampleToResolution(spacing=[1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], attributes) + Resample(spacing=[1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], attributes) message = str(excinfo.value) assert "case 'case'" in message and "3-dimensional grid" in message def test_resample_to_shape_transform_shape_dimension_mismatch_message(): - """ResampleToShape raises a formatted (f-string) message on a shape/target mismatch.""" + """A shape resample raises a formatted (f-string) message on a shape/target mismatch.""" attributes = Attribute() attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) with pytest.raises(TransformError) as excinfo: - ResampleToShape(shape=[4, 4]).transform_shape("group", "case", [10, 10, 10], attributes) + Resample(shape=[4, 4]).transform_shape("group", "case", [10, 10, 10], attributes) message = str(excinfo.value) assert "shape of 2 value(s)" in message assert "3 spatial axis/axes" in message @@ -332,7 +331,7 @@ def test_resample_to_shape_transform_shape_dimension_mismatch_message(): def test_resample_to_shape_does_not_mutate_config(): """#9 transform_shape must not write resolved dims back into the shared instance config.""" - resampler = ResampleToShape(shape=[0, 16, 16]) + resampler = Resample(shape=[0, 16, 16]) attributes = Attribute() attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) @@ -345,7 +344,7 @@ def test_resample_to_shape_does_not_mutate_config(): def test_resample_to_shape_inverse_without_spacing_metadata(): """Inverting a resample must not pop a 'Spacing' the forward pass never pushed.""" - resampler = ResampleToShape(shape=[4, 4, 4]) + resampler = Resample(shape=[4, 4, 4]) attributes = Attribute() # no image metadata at all tensor = torch.arange(8 * 8 * 8, dtype=torch.float32).reshape(1, 8, 8, 8) @@ -358,7 +357,7 @@ def test_resample_to_shape_inverse_without_spacing_metadata(): def test_resample_to_shape_inverse_pops_pushed_spacing(): """When 'Spacing' exists, the inverse removes the version the forward pass pushed.""" - resampler = ResampleToShape(shape=[4, 4, 4]) + resampler = Resample(shape=[4, 4, 4]) attributes = Attribute() attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) tensor = torch.zeros(1, 8, 8, 8) diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index bd317756..a0b441f1 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -63,10 +63,7 @@ PatchLocality, Percentage, Reduce, - ResampleToReference, - ResampleToResolution, - ResampleToShape, - ResampleTransform, + Resample, Save, SegmentationDisagreement, SelectLabel, @@ -76,7 +73,6 @@ Sum, Transform, Variance, - Warp, Write, ) from konfai.utils.dataset import Attribute, Dataset @@ -172,40 +168,33 @@ class _Case: "MergeLabels": [_Case(MergeLabels(), group="Ensemble")], "OneHot": [_Case(OneHot(4), group="Labels")], "Percentage": [_Case(Percentage(100.0))], - # The defaults ([1, 1, 1] mm / [100, 256, 256]) would be a no-op resample and a 6.5M-voxel upsample. - "ResampleToResolution": [ - _Case(ResampleToResolution([2.0, 1.0, 3.0])), # factorises: bit-identical - _Case(ResampleToResolution([2.0, 1.0, 3.0]), group="Int16", atol=_LSB_ATOL), + # The default (the case's own grid, no map) would be a no-op resample; these are the family's + # meaningful configurations, one per way of naming the grid and the map. + "Resample": [ + _Case(Resample(spacing=[2.0, 1.0, 3.0])), # factorises: bit-identical + _Case(Resample(spacing=[2.0, 1.0, 3.0]), group="Int16", atol=_LSB_ATOL), # uint8 resamples by nearest neighbour: no interpolation weights, so no rounding to disagree on. - _Case(ResampleToResolution([2.0, 1.0, 3.0]), group="Labels"), - ], - "ResampleToShape": [_Case(ResampleToShape([12, 8, 14]))], # factorises: bit-identical - # 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"), - # 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(Resample(spacing=[2.0, 1.0, 3.0]), group="Labels"), + _Case(Resample(shape=[12, 8, 14])), # factorises: bit-identical + # Onto a grid of its own, so part of the target reads from outside the case and takes the + # fill. atol is 0: both paths run the SAME sampler over global coordinates, so they agree + # bit for bit or not at all. + _Case(Resample(reference=_CASE_NAME, reference_group="Reference")), + _Case(Resample(reference=_CASE_NAME, reference_group="Reference"), group="Labels"), # Through a field the map does not factorise, so the blend goes to grid_sample -- which # normalises by the extent it is handed, and a patch is handed a window. That is the one # place a streamed answer is not bit-identical to the whole-volume one, and the atol says so. _Case( - ResampleToReference(entry=_CASE_NAME, group="Reference", field_group="Field"), + Resample(reference=_CASE_NAME, reference_group="Reference", field_group="Field"), atol=_REGRID_ATOL, ), + # A stored map never factorises: the grid_sample path again. + _Case(Resample(transforms={"transform": True}), atol=_REGRID_ATOL), + # A field on disk this registry cannot build: the streamed-equals-whole proof lives in + # test_warp.py, where a field exists -- including the measured, bound-less windows. + _Case(Resample(field="Dataset:h5", field_group="DVF"), sweep=False), ], - # A stored map never factorises, so this is the grid_sample path (see _REGRID_ATOL). - "ResampleTransform": [_Case(ResampleTransform({"transform": True}), atol=_REGRID_ATOL)], "Save": [_Case(Save("Dataset"))], - # Warp needs a field on disk to run, which this registry cannot build. Its - # streamed-equals-whole-volume proof lives in test_warp.py, where a field exists — including - # the bound-less case, whose windows are measured from the field at run. - "Warp": [_Case(Warp(field="Dataset:h5", group="DVF"), sweep=False)], # Reduce is a cardinality marker the cohort engine splits out of the chain, never a per-case # stage: it declares WHOLE_VOLUME so a chain reaching the ordinary planner refuses rather than # streams, which is what puts it out of the equivalence sweep below. @@ -720,7 +709,7 @@ def manager(transform: Transform) -> DatasetManager: data_augmentations_list=[augmentations], ) - resample = ResampleToResolution([2.0, 1.0, 3.0]) + resample = Resample([2.0, 1.0, 3.0]) assert manager(resample).can_stream_patch(1) is True # Two regions in one chain — the resample's, then the flip draw's — composed into one pull. flip = FlipAugmentation(f_prob=[1.0, 1.0, 1.0]) diff --git a/tests/unit/test_transformer_workflow.py b/tests/unit/test_transformer_workflow.py index 1593418d..0a1a92de 100644 --- a/tests/unit/test_transformer_workflow.py +++ b/tests/unit/test_transformer_workflow.py @@ -128,7 +128,7 @@ def test_streamable_chain_plans_streams_and_writes(tmp_path: Path) -> None: _REGION_CHAIN = """\ - ResampleToResolution: + Resample: spacing: [1.0, 1.0, 1.0] Write: dataset: {out}:h5 @@ -399,9 +399,9 @@ def spy(destination, group, shape, dtype, attributes): Resample: spacing: [2.0, 2.0, 2.0] align: origin - ResampleToReference: - entry: CASE_000 - group: CT + konfai.data.transform:Resample: + reference: CASE_000 + reference_group: CT Write: dataset: {out}:h5 """ @@ -414,10 +414,7 @@ def test_a_stage_is_asked_about_its_own_input_not_the_case_as_stored(tmp_path: P edge falls short of the reference's and part of the output will be fill. Asked about the case as STORED it covers all of it and says nothing -- and the plan would stay silent about that fill. - (With ``align: extent``, the default, the box is preserved and the answer is honestly 100%. - 67.5% is also exactly what the old ``ResampleToResolution`` recorded here -- because its header - said origin-aligned while its data was extent-aligned, which is the mismatch this stage no - longer has.) + (With ``align: extent``, the default, the box is preserved and the answer is honestly 100%.) """ _write_source(tmp_path) _write_config(tmp_path, _RESAMPLED_THEN_REFERENCED.format(out=tmp_path / "out")) @@ -493,7 +490,7 @@ def test_overwrite_rewrites_a_geometry_changing_chain_correctly(tmp_path: Path) _write_source(tmp_path) _write_config( tmp_path, - " ResampleToResolution:\n" + " Resample:\n" " spacing: [1.0, 3.0, 4.0]\n" " Write:\n" f" dataset: {tmp_path / 'out'}:h5\n", diff --git a/tests/unit/test_warp.py b/tests/unit/test_warp.py index ddb48af4..9bf9edd1 100644 --- a/tests/unit/test_warp.py +++ b/tests/unit/test_warp.py @@ -14,7 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""``Warp`` resamples a case through a displacement field, region by region. +"""``Resample`` through a displacement field alone: a warp on the case's own grid, region by region. The claim under test is the one that matters for a volume larger than memory: the streamed result equals the whole-volume one, each region's window is sized from the field values it reads, and a @@ -26,7 +26,7 @@ import pytest import torch from konfai.data.patching import DatasetManager -from konfai.data.transform import LocalityKind, RegionContext, Save, Warp +from konfai.data.transform import LocalityKind, RegionContext, Resample, Save from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, Attribute, Dataset from konfai.utils.errors import TransformError from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE, _zarr_v3_available @@ -82,7 +82,7 @@ def _manager(source: Dataset, transforms: list) -> DatasetManager: ) -def _recorded(warp: Warp, attribute: Attribute | None = None, shape: tuple[int, ...] = (10, 12, 14)) -> Warp: +def _recorded(warp: Resample, attribute: Attribute | None = None, shape: tuple[int, ...] = (10, 12, 14)) -> Resample: """A stage that has met its case — which is when a region can be asked about at all.""" warp.transform_shape("CT", "CASE_000", list(shape), attribute if attribute is not None else _attributes()) return warp @@ -97,7 +97,7 @@ def test_the_source_region_is_the_target_grown_by_the_field_reach(tmp_path: Path cannot express at all. """ _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 4.0, 4.0)) - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) + warp = _recorded(Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF")) assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID target = (slice(4, 6), slice(4, 6), slice(4, 6)) @@ -117,9 +117,9 @@ def test_the_source_region_is_the_target_grown_by_the_field_reach(tmp_path: Path def test_an_oblique_case_grows_its_window_on_every_axis(tmp_path: Path) -> None: """The bug a per-axis halo hid: a displacement along x reaches into y and z when the axes turn. - ``Warp`` used to convert a world bound to a halo per ARRAY axis, which silently assumed the - direction cosines were the identity -- on a turned case the window was short on the axes the - displacement actually reached, and a short window returns the border value rather than raising. + A world bound converted to a halo per ARRAY axis silently assumes the direction cosines are + the identity -- on a turned case the window is short on the axes the displacement actually + reaches, and a short window returns the border value rather than raising. """ turned = _attributes() angle = np.deg2rad(35.0) @@ -127,7 +127,7 @@ def test_an_oblique_case_grows_its_window_on_every_axis(tmp_path: Path) -> None: turned["Direction"] = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]).reshape(-1) _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 0.0, 0.0)) - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF"), turned) + warp = _recorded(Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF"), turned) target = (slice(5, 6), slice(5, 6), slice(5, 6)) window = warp.measured_region_source("CASE_000", target, [10, 12, 14], turned) @@ -149,7 +149,7 @@ def test_the_bound_the_fields_recorded_prices_the_plan(tmp_path: Path) -> None: attribute[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" fields.write("DVF", case, field, attribute) - warp = Warp(field=f"{tmp_path / 'dvf'}:omezarr", group="DVF") + warp = Resample(field=f"{tmp_path / 'dvf'}:omezarr", field_group="DVF") locality = warp.patch_locality(_attributes()) # The cohort's bound is (x=1.0, y=6.0, z=3.0); spacing in array order (z, y, x) is (1, 1, 2), so @@ -176,7 +176,7 @@ def test_the_header_scan_survives_an_unreadable_entry_in_the_field_group(tmp_pat # The first entry the scan meets, so the read reaches it before any bound-less entry ends the scan. (tmp_path / "dvf" / "CASE_000" / "DVF.mha").write_bytes(b"not an image") - locality = Warp(field=f"{tmp_path / 'dvf'}:mha", group="DVF").patch_locality(_attributes()) + locality = Resample(field=f"{tmp_path / 'dvf'}:mha", field_group="DVF").patch_locality(_attributes()) assert locality.kind is LocalityKind.WHOLE_VOLUME @@ -186,7 +186,7 @@ def test_a_field_with_no_bound_still_streams_with_windows_measured_at_run(tmp_pa for sampling regardless, and the sup of those very values sizes that region's source pull — per region, so a quiet slab pays a quiet halo where the shifted one pays its shift.""" _source, _fields, _volume = _fixture(tmp_path, shift_um=(4.0, 0.0, 0.0)) # 4 um along x alone - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) + warp = _recorded(Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF")) assert warp.patch_locality(_attributes()).kind is LocalityKind.REGRID @@ -203,7 +203,7 @@ def test_a_field_with_no_bound_still_streams_with_windows_measured_at_run(tmp_pa def test_sizing_and_sampling_share_one_field_read(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """The window that sizes a region's pull is the window the sampler needs next: one read.""" _source, _fields, _volume = _fixture(tmp_path) - warp = _recorded(Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF")) + warp = _recorded(Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF")) displacement = warp.displacement assert displacement is not None reads: list[int] = [] @@ -222,7 +222,7 @@ def test_sizing_and_sampling_share_one_field_read(tmp_path: Path, monkeypatch: p def test_a_constant_shift_moves_the_volume_by_that_many_voxels(tmp_path: Path) -> None: _source, _fields, volume = _fixture(tmp_path, shift_um=(0.0, 0.0, 3.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") + warp = Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF") moved = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() @@ -237,7 +237,7 @@ def test_streamed_equals_whole_volume(tmp_path: Path, monkeypatch: pytest.Monkey monkeypatch.setattr(patching_module, "_SWEEP_SLAB_ROWS", 3) source, _fields, volume = _fixture(tmp_path, shift_um=(1.0, 2.0, 3.0)) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") + warp = Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF") reference = warp("CASE_000", torch.from_numpy(volume), _attributes()).numpy() @@ -262,7 +262,7 @@ def test_a_field_beyond_its_recorded_bound_raises(tmp_path: Path) -> None: field = np.zeros((3, 10, 12, 14), dtype=np.float32) field[2] = 9.0 Dataset(tmp_path / "lying", "mha").write("DVF", "CASE_000", field, attributes) - warp = Warp(field=f"{tmp_path / 'lying'}:mha", group="DVF") + warp = Resample(field=f"{tmp_path / 'lying'}:mha", field_group="DVF") _recorded(warp) with pytest.raises(TransformError, match=r"on component 2, beyond the 1\.000"): @@ -279,17 +279,17 @@ def test_a_field_with_the_wrong_component_count_is_named(tmp_path: Path) -> None rng = np.random.default_rng(1) Dataset(tmp_path / "src", "h5").write("CT", "CASE_000", rng.random((1, 4, 4, 4)).astype(np.float32), _attributes()) Dataset(tmp_path / "dvf", "h5").write("DVF", "CASE_000", np.zeros((2, 4, 4, 4), np.float32), _attributes()) - warp = Warp(field=f"{tmp_path / 'dvf'}:h5", group="DVF") + warp = Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF") with pytest.raises(TransformError, match="component"): warp("CASE_000", torch.zeros(1, 4, 4, 4), _attributes()) -def test_warp_without_a_field_is_refused_at_construction() -> None: - with pytest.raises(TransformError, match="needs a 'field'"): - Warp(field="") +def test_an_empty_field_path_declares_no_field() -> None: + """An empty ``field`` with no group is the identity map, not a broken declaration.""" + assert Resample(field="").displacement is None def test_unknown_interpolation_is_refused_at_construction() -> None: with pytest.raises(TransformError, match="unknown interpolation"): - Warp(field="./x:h5", interpolation="cubic") + Resample(field="./x:h5", interpolation="cubic") diff --git a/tests/unit/test_write_pyramid_and_field_bound.py b/tests/unit/test_write_pyramid_and_field_bound.py index d1f6a6a5..7bd3716e 100644 --- a/tests/unit/test_write_pyramid_and_field_bound.py +++ b/tests/unit/test_write_pyramid_and_field_bound.py @@ -210,13 +210,13 @@ def test_the_recorded_bound_reaches_each_axis_by_its_own_spacing_under_anisotrop wrong neighbourhood -- and under this anisotropy the three numbers are far enough apart (3, 22 and 31 voxels) that any permutation of them is visible. """ - from konfai.data.transform import LocalityKind, Warp + from konfai.data.transform import LocalityKind, Resample field = np.zeros((3, 8, 8, 8), dtype=np.float32) field[0, 4, 4, 4], field[1, 2, 2, 2], field[2, 1, 1, 1] = 917.5, -640.25, 96.0 store = tmp_field_store(field) - warp = Warp(field=f"{store}:omezarr", group="DVF") + warp = Resample(field=f"{store}:omezarr", field_group="DVF") attribute = Attribute() attribute["Spacing"] = np.array([30.08, 30.08, 40.0]) # stored (x, y, z) attribute["Origin"] = np.zeros(3) From d0ae5abc55fb7d9afaef3f1dc30b954668411fde Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 02:49:45 +0200 Subject: [PATCH 23/39] fix(impact-reg): one staging root per group, so mixed forms stay readable A directory dataset's backend is detected from its first case, and one store entry flips the whole root to the store backend: an OME-Zarr moving staged beside the .mha field every published preset declares made the field unreadable at the first register. Each group now stages into a root of its own -- homogeneous, so each keeps its backend -- and the run reads its roots side by side. Also: re-staging a case replaces its link instead of raising, and _the_output names the stale other-form file it found and the remedy. --- .../impact_reg_konfai/impact_reg.py | 39 ++++++++------- .../tests/unit/test_orchestration.py | 47 +++++++++++++++++++ 2 files changed, 66 insertions(+), 20 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 5ff2d545..779cef2d 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -212,7 +212,7 @@ def _run_transform( api.transform( name, - f"{root}:mha", # the format token is only the FIRST read candidate; entries resolve by name + datasets, chains, gpu=list(gpu), cpu=cpu or 1, @@ -442,14 +442,14 @@ def _ensemble_mean( """Average one case's preset fields into ``//DVF`` — Reduce(Mean), streamed.""" from konfai.data.transform import Reduce, Write - root = work / f"ensemble_{case}" - for preset, dvf in zip(presets, dvf_paths, strict=True): - _stage(root, preset, "DVF", dvf) + members = _stage_group( + work / f"ensemble_{case}", "DVF", {preset: dvf for preset, dvf in zip(presets, dvf_paths, strict=True)} + ) suffixes = "".join(dvf_paths[0].suffixes) _output_path(output / case, "DVF", suffixes) # drop a stale other-form DVF before writing _run_transform( f"impact_reg_ensemble_{case}", - root, + [members], { "DVF": { "DVF": [ @@ -483,17 +483,16 @@ def _derive_moved( """ from konfai.data.transform import Resample, Write - root = work / "moved_stage" + fields = {case: _the_output(output / case, "DVF") for case in cases} suffixes = "" - for case, moving in cases.items(): - dvf = _the_output(output / case, "DVF") - _stage(root, case, "Moving", moving) - _stage(root, case, "DVF", dvf) + for case, dvf in fields.items(): suffixes = "".join(dvf.suffixes) _output_path(output / case, "Moved", suffixes) # drop a stale other-form Moved + moving_root = _stage_group(work / "moved_stage", "Moving", cases) + field_root = _stage_group(work / "moved_stage", "DVF", fields) _run_transform( "impact_reg_moved", - root, + [moving_root, field_root], { "Moving": { "Moved": [ @@ -626,19 +625,21 @@ def _warp_onto_fixed( """ from konfai.data.transform import Resample, Write - root = work / f"warp_{kind}" - _stage(root, "P000", "Fixed", fixed) - _stage(root, "P000", "Moving", moving) + base = work / f"warp_{kind}" + datasets = [ + _stage_group(base, "Fixed", {"P000": fixed}), + _stage_group(base, "Moving", {"P000": moving}), + ] resample: dict[str, object] = {"reference": "{case}", "reference_group": "Fixed"} if kind == "seg": resample["interpolation"] = "nearest" if transform_path is not None: - _stage(root, "P000", "Reg", transform_path) + datasets.append(_stage_group(base, "Reg", {"P000": transform_path})) resample["transforms"] = {"Reg": False} out_root = work / f"moved_{kind}" _run_transform( f"impact_reg_eval_{kind}", - root, + datasets, {"Moving": {"Moved": [Resample(**resample), Write(dataset=f"{out_root}:mha")]}}, work, gpu, @@ -673,14 +674,12 @@ def uncertainty( try: from konfai.data.transform import Magnitude, Reduce, Write - root = work / "members" members = _units(list(dvfs)) - for index, dvf in enumerate(members): - _stage(root, f"M{index:03d}", "DVF", dvf) + spec = _stage_group(work / "members", "DVF", {f"M{index:03d}": dvf for index, dvf in enumerate(members)}) suffixes = "".join(members[0].suffixes) _run_transform( "impact_reg_uncertainty", - root, + [spec], { "DVF": { "Uncertainty": [ diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index 94bb24ed..e7a147ff 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -209,3 +209,50 @@ def test_register_fields_only_writes_nothing_derived(tmp_path: Path) -> None: assert (case / "DVF.mha").is_file() assert not (case / "Moved.mha").exists() assert not (case / "Transform.h5").exists() + + +def test_register_reads_a_store_moving_against_an_itk_field(tmp_path: Path) -> None: + """One store entry beside an ``.mha`` flips a mixed root's backend — the staging keeps one root + per group, so a caller's OME-Zarr moving registers against the ``.mha`` field every published + preset declares.""" + ome_zarr = pytest.importorskip("konfai.utils.ome_zarr") + if not ome_zarr._zarr_v3_available(): + pytest.skip("writing an OME-Zarr moving needs zarr 3") + + volume = np.arange(8**3, dtype=np.float32).reshape(8, 8, 8) + moving = tmp_path / "moving.ome.zarr" + ome_zarr.write_ome_zarr(moving, volume[None], spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + reference = sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)) # the moving's grid + app = reg.ImpactRegKonfAIApp() + + def field_only(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + out = Path(work) / preset / "P000" + out.mkdir(parents=True, exist_ok=True) + _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) + return {"P000": out / "DVF.mha"} + + app._infer_preset = field_only # type: ignore[method-assign] + out = tmp_path / "Output" + app.register(["FireANTs_SyN"], [fixed], [moving], output=out) + + # moved(p) = moving(p + d), d = +2 along x on a unit grid: moving is z*64 + y*8 + x. + moved = sitk.GetArrayFromImage(sitk.ReadImage(str(out / "P000" / "Moved.mha"))) + np.testing.assert_allclose(moved[0, 0, 0], 2.0, atol=1e-6) + np.testing.assert_allclose(moved[1, 1, 0], 74.0, atol=1e-6) + assert (out / "P000" / "Transform.h5").is_file() + + +def test_stage_group_replaces_an_existing_link(tmp_path: Path) -> None: + """Re-staging the same case points the link at the new source instead of raising.""" + first, second = tmp_path / "a.mha", tmp_path / "b.mha" + for path in (first, second): + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((2, 2, 2), dtype=np.float32)), str(path)) + + reg._stage_group(tmp_path / "stage", "DVF", {"P000": first}) + spec = reg._stage_group(tmp_path / "stage", "DVF", {"P000": second}) + + root = Path(spec.rpartition(":")[0]) + assert (root / "P000" / "DVF.mha").resolve() == second.resolve() From 3a2458c4e97a653b79a098bef7efacfbbf9c1760 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 02:50:00 +0200 Subject: [PATCH 24/39] perf(impact-reg): fill Transform.h5 region by region sitk.WriteTransform needs the whole field resident, converted to float64 -- the run's peak, per case, once everything else streams. The constraint is the function, not the format: an ITK transform file is three HDF5 datasets (type; size/origin/spacing/direction; the field buffer, component fastest), and HDF5 writes by regions. The file is now written through h5py with the parameters streamed from the store, so the peak is one slab in float64; without h5py the sitk path serves, whole. The test pins read-back equality with sitk's own writer -- same type, fixed parameters and parameters, exactly, from either field form. --- .../impact_reg_konfai/impact_reg.py | 123 +++++++++++++++--- .../tests/unit/test_displacement_field_io.py | 19 +++ .../tests/unit/test_orchestration.py | 4 +- 3 files changed, 126 insertions(+), 20 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 779cef2d..b9ff3aed 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -173,29 +173,112 @@ def _copy_output(src: Path, dest_dir: Path, stem: str) -> Path: return dest -def _stage(root: Path, case: str, group: str, source: Path) -> None: - """Link one cohort entry — ``root//`` → ``source``. No bytes move. - - A dataset is a root of CASES holding GROUPS; paths from the command line (or another run's - output tree) become one by symlink, exactly as konfai-apps stages its own inputs. Entries - resolve by name whatever their form — an ITK file, an OME-Zarr store, a stored transform. +def _stage_group(base: Path, group: str, entries: dict[str, Path]) -> str: + """One dataset ROOT holding one GROUP — ``base///`` symlinks — + returned as the ``path:format`` spec a run consumes. No bytes move. + + One root per group, not one mixed cohort: a directory dataset's backend is detected from its + first case, and a single store entry flips the whole root to the store backend — the ``.mha`` + beside it stops resolving. A homogeneous root keeps every entry readable whatever mix of forms + the caller and the presets produced; the run reads all its roots side by side. """ - case_dir = root / case - case_dir.mkdir(parents=True, exist_ok=True) - (case_dir / (group + "".join(source.suffixes))).symlink_to(source.resolve()) + forms = {"".join(source.suffixes).lower() for source in entries.values()} + if len(forms) > 1: + raise RuntimeError( + f"group '{group}' mixes storage forms ({', '.join(sorted(forms))}); a directory dataset" + " has one backend, so a mixed group cannot be staged. Re-run the producers to one form." + ) + root = base / group + suffixes = "" + for case, source in entries.items(): + case_dir = root / case + case_dir.mkdir(parents=True, exist_ok=True) + suffixes = "".join(source.suffixes) + # Clear every form of the entry, not only the current one: a re-stage that switched forms + # would otherwise leave two links and discovery by stem finds both. + for stale in case_dir.glob(f"{group}.*"): + stale.unlink() if stale.is_symlink() or stale.is_file() else shutil.rmtree(stale) + link = case_dir / (group + suffixes) + link.symlink_to(source.resolve()) + return f"{root}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.') or 'mha')}" def _the_output(dest_dir: Path, stem: str) -> Path: """The single output named ``stem`` in ``dest_dir`` — one, exactly.""" matches = sorted(dest_dir.glob(f"{stem}.*")) if len(matches) != 1: - raise FileNotFoundError(f"Expected exactly one '{stem}' under {dest_dir}, found {len(matches)}.") + found = ", ".join(path.name for path in matches) or "none" + raise FileNotFoundError( + f"Expected exactly one '{stem}' under {dest_dir}, found {found}. More than one is a" + " stale other-form output left beside the current one: remove it, or re-run register," + " which clears the stem before writing." + ) return matches[0] +def _write_displacement_transform(dvf: Path, dest: Path, work: Path) -> None: + """``dest`` as ITK's HDF5 transform writer lays it out — parameters filled region by region. + + ``sitk.WriteTransform`` needs the whole field resident, in float64; the FORMAT does not: an + ITK transform file is three HDF5 datasets (the type; the fixed parameters — size, origin, + spacing, direction; the parameters — the field buffer, component fastest), and HDF5 writes by + regions. The field is read in slabs through ``Dataset``, so the peak is one slab in float64 + instead of the field twice over. Read-back equality with sitk's own writer is pinned by the + app tests. Without ``h5py`` the sitk path serves, whole. + """ + try: + import h5py + except ImportError: + sitk.WriteTransform(sitk.DisplacementFieldTransform(read_displacement_field(dvf)), str(dest)) + return + from konfai.utils.dataset import Dataset + + spec = _stage_group(work / "transform_h5", "DVF", {"P000": dvf}) + filename, _colon, file_format = spec.rpartition(":") + dataset = Dataset(Path(filename), file_format) + shape, attribute = dataset.get_infos("DVF", "P000") + channels, spatial = int(shape[0]), [int(extent) for extent in shape[1:]] + if channels != 3 or len(spatial) != 3: + raise RuntimeError(f"Transform.h5 needs a 3-component 3-D field; '{dvf}' has shape {list(shape)}.") + fixed = np.concatenate( + [ + np.asarray(spatial[::-1], dtype=np.float64), # size, in (x, y, z) + attribute.get_np_array("Origin").astype(np.float64), + attribute.get_np_array("Spacing").astype(np.float64), + attribute.get_np_array("Direction").astype(np.float64).reshape(-1), + ] + ) + rows, offset = 16, 0 + streamed = dataset.bounded_region_reads("DVF", "P000") + resident = None if streamed else np.asarray(dataset.read_data("DVF", "P000")[0]) + staging = dest.with_name(dest.name + ".tmp") + with h5py.File(staging, "w") as file: + file.create_dataset( + "TransformGroup/0/TransformType", + data=[b"DisplacementFieldTransform_double_3_3"], + dtype=h5py.string_dtype(encoding="ascii"), # ITK reads a variable-length ASCII string + ) + file.create_dataset("TransformGroup/0/TransformFixedParameters", data=fixed) + parameters = file.create_dataset( + "TransformGroup/0/TransformParameters", shape=(3 * int(np.prod(spatial)),), dtype=np.float64 + ) + for start in range(0, spatial[0], rows): + stop = min(start + rows, spatial[0]) + if resident is None: + block, _ = dataset.read_data_slice( + "DVF", "P000", (slice(None), slice(start, stop), slice(None), slice(None)) + ) + else: + block = resident[:, start:stop] + slab = np.moveaxis(np.asarray(block, dtype=np.float64), 0, -1).ravel() + parameters[offset : offset + slab.size] = slab + offset += slab.size + os.replace(staging, dest) + + def _run_transform( name: str, - root: Path, + datasets: list[str], chains: dict, work: Path, gpu: list[int], @@ -352,6 +435,13 @@ def register( # volume belongs to which case, in the same order it will use, so a dataset in and a single pair # in go down one path. moving_units = _units(moving_images) + for side, masks in (("fixed", fixed_masks), ("moving", moving_masks)): + mask_units = _units(list(masks)) if masks else [] + if masks and len(mask_units) != len(moving_units): + raise RuntimeError( + f"the {side} masks expand to {len(mask_units)} unit(s) for {len(moving_units)}" + " case(s); masks pair with cases by position, so the counts must match." + ) work = _work_dir(tmp_dir, "impact_reg_") try: @@ -417,14 +507,11 @@ def register( # map -- one interpolation, slab by slab, the whole cohort under one plan. self._derive_moved(dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet) for case in cases: - # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid field - # as a SimpleITK transform. Inherently whole -- the .h5 format carries the full - # field -- which is why it goes with the moved image rather than being - # unconditional. - transform = sitk.DisplacementFieldTransform( - read_displacement_field(_the_output(output / case, "DVF")) + # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid + # field as an ITK transform file, filled region by region from the DVF. + _write_displacement_transform( + _the_output(output / case, "DVF"), output / case / "Transform.h5", work ) - sitk.WriteTransform(transform, str(output / case / "Transform.h5")) finally: shutil.rmtree(work, ignore_errors=True) diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index d8efdb30..1194af62 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -258,3 +258,22 @@ def test_transform_reads_back_identically_from_either_form(tmp_path: Path, suffi reference = sitk.DisplacementFieldTransform(sitk.Image(original)) for point in ((9.0, -1.0, 12.0), (7.5, -2.5, 11.0)): assert restored.TransformPoint(point) == pytest.approx(reference.TransformPoint(point)) + + +@pytest.mark.parametrize("suffix", [".mha", ".ome.zarr"]) +def test_transform_h5_reads_back_as_sitk_would_have_written_it(tmp_path: Path, suffix: str) -> None: + """The streamed writer and ``sitk.WriteTransform`` must be the same file to ITK's reader: + same type, same fixed parameters (size, origin, spacing, direction), same parameters — exactly.""" + from impact_reg_konfai.impact_reg import _write_displacement_transform + + original = _field() + _write_displacement_field(original, tmp_path / f"DVF{suffix}") + work = tmp_path / "work" + work.mkdir() + + _write_displacement_transform(tmp_path / f"DVF{suffix}", tmp_path / "Transform.h5", work) + + got = sitk.ReadTransform(str(tmp_path / "Transform.h5")) + want = sitk.DisplacementFieldTransform(read_displacement_field(tmp_path / f"DVF{suffix}")) + assert got.GetFixedParameters() == want.GetFixedParameters() + assert got.GetParameters() == want.GetParameters() diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index e7a147ff..8343cfad 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -54,8 +54,8 @@ def test_get_available_presets_keeps_only_registration_apps(tmp_path: Path, monk def test_find_outputs_raises_when_missing(tmp_path: Path) -> None: - with pytest.raises(FileNotFoundError, match=r"Moved\.mha"): - reg._find_outputs(tmp_path, "Moved.mha") + with pytest.raises(FileNotFoundError, match="Moved"): + reg._find_outputs(tmp_path, "Moved") # --------------------------------------------------------------------------- mask sentinel From b7a39f2f6d9374fd416fd43629c2fd42f5e44a01 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 03:09:44 +0200 Subject: [PATCH 25/39] docs: migrate the gallery and say the old Resample names are gone The changelog told readers 'the old names still work and are thin argument translations'. That was true of the commit that made them forwarding shells, not of the one that removed them -- and it is the section the CI publishes as the release note, so it would have sent people migrating after the fact, onto something broken. It now carries the argument mapping. The visual gallery still imported ResampleToShape and ResampleToResolution. --- CHANGELOG.md | 4 +++- docs/scripts/generate_visual_gallery.py | 7 +++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9cdd517..ad575c0e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,7 +26,9 @@ written replaces it. **which grid to write on** (nothing, `spacing`, `shape`, or `reference` — a stored image's grid adopted whole) and **what map to write it through** (`field`, `transforms`, or neither). Every combination is legal, and asked for together they compose into **one interpolation** instead of - two. The old names still work and are thin argument translations. + two. **The old names are gone**, so a config naming one must be migrated: the class becomes + `Resample`, and `ResampleToReference`'s `entry` / `group` / `dataset` become `reference` / + `reference_group` / `reference_dataset`. Every other argument carries over unchanged. - **transform**: `align` says where a `spacing` or `shape` grid sits — `extent` (the default) keeps the field of view, `origin` keeps voxel zero's centre. This was decided silently before, and differently by the data and by the header. diff --git a/docs/scripts/generate_visual_gallery.py b/docs/scripts/generate_visual_gallery.py index 3b0b1b18..7d7fb716 100644 --- a/docs/scripts/generate_visual_gallery.py +++ b/docs/scripts/generate_visual_gallery.py @@ -31,8 +31,7 @@ Normalize, Padding, Permute, - ResampleToResolution, - ResampleToShape, + Resample, Standardize, ) from konfai.utils.dataset import Attribute @@ -142,12 +141,12 @@ def main() -> None: target_shape = [220, 220] shape_attribute = Attribute() shape_attribute["Spacing"] = np.asarray([0.8, 0.8]) - resampled_shape = ResampleToShape(shape=target_shape)("IMAGE", normalized.clone(), shape_attribute) + resampled_shape = Resample(shape=target_shape)("IMAGE", normalized.clone(), shape_attribute) resolution_attribute = Attribute() resolution_attribute["Spacing"] = np.asarray([0.8, 0.8]) target_spacing = [1.25, 0.55] - resampled_resolution = ResampleToResolution(spacing=target_spacing, inverse=True)( + resampled_resolution = Resample(spacing=target_spacing, inverse=True)( "IMAGE", normalized.clone(), resolution_attribute ) From b0ff7f2a5bb50b125bdd49fd4973d8c8d5d6ce50 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 03:09:45 +0200 Subject: [PATCH 26/39] feat(data): an ITK transform file is a Dataset backend The ':itktransform' format writes a displacement field as the ITK transform file any ITK consumer loads -- three HDF5 datasets (the type, ASCII as ITK reads it; size/origin/spacing/direction; the field buffer, component fastest) -- and fills the parameters region by region: sitk.WriteTransform needs the whole field resident in float64 where the file itself writes by slabs. Whole and streamed writes are pinned identical to sitk's own writer through ITK's reader; an aborted stream leaves no entry under the final name. The read side hands back what Dataset.read_transform decodes -- a displacement entry with its field and marker, any other stored transform as parameter rows -- so a staged .h5/.tfm resolves through the same Dataset surface as every other entry. Without h5py the sitk path serves, whole. --- konfai/utils/dataset.py | 217 +++++++++++++++++++++++ tests/unit/test_itk_transform_backend.py | 124 +++++++++++++ 2 files changed, 341 insertions(+) create mode 100644 tests/unit/test_itk_transform_backend.py diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 158623e2..97a2c39e 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -721,6 +721,75 @@ def _close(self, success: bool) -> None: parent.move(temporary_name, self._final_name) +def _create_itk_transform_file(path: str, spatial: list[int], attributes: Attribute) -> tuple[Any, Any]: + """An ITK displacement-transform HDF5 file with its parameters dataset still to fill. + + Three datasets, as ITK's own writer lays them out: the type (a variable-length ASCII string, + which is what ITK's reader accepts), the fixed parameters — size, origin, spacing, direction — + and the parameters, the field buffer with the component fastest, float64. Returns the open file + and the parameters dataset. + """ + import h5py + + fixed = np.concatenate( + [ + np.asarray(spatial[::-1], dtype=np.float64), # size, in (x, y, z) + attributes.get_np_array("Origin").astype(np.float64), + attributes.get_np_array("Spacing").astype(np.float64), + attributes.get_np_array("Direction").astype(np.float64).reshape(-1), + ] + ) + file = h5py.File(path, "w") + file.create_dataset( + "TransformGroup/0/TransformType", + data=[b"DisplacementFieldTransform_double_3_3"], + dtype=h5py.string_dtype(encoding="ascii"), + ) + file.create_dataset("TransformGroup/0/TransformFixedParameters", data=fixed) + parameters = file.create_dataset( + "TransformGroup/0/TransformParameters", shape=(3 * int(np.prod(spatial)),), dtype=np.float64 + ) + return file, parameters + + +class _ItkTransformDataStream(DataStream): + """An ITK displacement-transform file written region by region. + + A slab of the field maps to one contiguous span of the parameters (the buffer is ``[z][y][x]`` + with the component fastest), so full-width leading-axis slabs — what the streamed write + dispatcher emits — land with plain offset writes. Under a temporary name until the clean exit, + like every stream. + """ + + def __init__(self, file: Any, parameters: Any, temporary_path: str, final_path: str, spatial: list[int]) -> None: + self._h5 = file + self._parameters = parameters + self._temporary_path = temporary_path + self._final_path = final_path + self._spatial = [int(extent) for extent in spatial] + + def write_slice(self, slices: tuple[slice, ...], data: np.ndarray) -> None: + channels, leading, *rest = slices + full = (channels.start or 0) == 0 and channels.stop in (None, 3) + for axis, part in enumerate(rest, start=2): + full = full and (part.start or 0) == 0 and part.stop in (None, self._spatial[axis]) + if not full: + raise DatasetManagerError( + "A transform file writes full-width leading-axis slabs, and this region is not one.", + "This is a bug if it was reached: the streamed write dispatcher finalizes full rows.", + ) + block = np.moveaxis(np.asarray(data, dtype=np.float64), 0, -1).ravel() + offset = 3 * int(leading.start or 0) * int(np.prod(self._spatial[2:], dtype=np.int64)) + self._parameters[offset : offset + block.size] = block + + def _close(self, success: bool) -> None: + self._h5.close() + if success: + os.replace(self._temporary_path, self._final_path) + else: + Path(self._temporary_path).unlink(missing_ok=True) + + # MetaImage ElementType for each NumPy dtype a streamed .mha can hold. _MHA_ELEMENT_TYPES = { "int8": "MET_CHAR", @@ -1866,6 +1935,150 @@ def get_infos(self, group: str, name: str) -> tuple[list[int], Attribute]: info = get_dicom_info(self._path(name)) return info["shape"], self._attributes(info) + class ItkTransformFile(AbstractFile): + """ITK transform files, one ``/.h5`` per entry. + + The write side is the point: ``sitk.WriteTransform`` needs the whole field resident in + float64, where the FILE is three HDF5 datasets that write by regions — so a displacement + field streams into a transform any ITK consumer (Slicer first) loads. The read side hands + back what ``Dataset.read_transform`` decodes: a displacement entry carries its field and + the displacement marker; any other stored transform, the parameter rows and type keys of + ``_encode_transform_leaves``. + """ + + def __init__(self, filename: str, read: bool) -> None: + self.filename = filename + self.read = read + + def __enter__(self): + return self + + def __exit__(self, exc_type, value, traceback): + pass + + def _path(self, name: str) -> str: + for extension in ("h5", "tfm"): + candidate = f"{self.filename}{name}.{extension}" + if os.path.exists(candidate): + return candidate + return f"{self.filename}{name}.h5" + + def file_to_data(self, group: str, name: str) -> tuple[np.ndarray, Attribute]: + transform = sitk.ReadTransform(self._path(name)) + attributes = Attribute() + if "DisplacementFieldTransform" in transform.GetName(): + field = sitk.DisplacementFieldTransform(transform).GetDisplacementField() + data, attributes = image_to_data(field) + attributes[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" + return data, attributes + leaves = _encode_transform_leaves(transform, name, attributes) + longest = max(len(leaf) for leaf in leaves) + return ( + np.asarray([np.pad(leaf, (0, longest - len(leaf)), constant_values=np.nan) for leaf in leaves]), + attributes, + ) + + def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: + data, attributes = self.file_to_data(group, name) + return data[slices], attributes + + def file_to_data_statistics( + self, + group: str, + name: str, + channels: list[int] | None = None, + ) -> dict[str, Any]: + data, _attributes = self.file_to_data(group, name) + if channels is not None: + data = data[channels] + return _finalize_running_statistics(_update_running_statistics(None, data)) + + def data_to_file( + self, + name: str, + data: sitk.Image | sitk.Transform | np.ndarray, + attributes: Attribute | None = None, + ) -> None: + os.makedirs(self.filename, exist_ok=True) + final = self._path(name) + staging = f"{self.filename}.{name}.{os.getpid()}.tmp.h5" + if isinstance(data, sitk.Transform): + sitk.WriteTransform(data, staging) + os.replace(staging, final) + return + if isinstance(data, sitk.Image): + data, attributes = image_to_data(data) + array = np.asarray(data) + if attributes is None or array.ndim != 4 or array.shape[0] != 3: + raise DatasetManagerError( + f"An ':itktransform' entry is a 3-component 3-D displacement field; '{name}' has" + f" shape {list(array.shape)}.", + "Write the field itself (channel-first, with its geometry), or a sitk.Transform.", + ) + try: + import h5py # noqa: F401 + except ImportError: + field = sitk.Cast(data_to_image(array, attributes), sitk.sitkVectorFloat64) + sitk.WriteTransform(sitk.DisplacementFieldTransform(field), staging) + os.replace(staging, final) + return + spatial = [int(extent) for extent in array.shape[1:]] + file, parameters = _create_itk_transform_file(staging, spatial, attributes) + with file: + parameters[:] = np.moveaxis(array.astype(np.float64), 0, -1).ravel() + os.replace(staging, final) + + def open_data_stream( + self, + name: str, + shape: list[int], + dtype: np.dtype, + attributes: Attribute, + region_shape: list[int] | None = None, + ) -> DataStream | None: + del dtype, region_shape # the parameters are float64 whatever arrives, converted per slab + try: + import h5py # noqa: F401 + except ImportError: + return None + if len(shape) != 4 or shape[0] != 3 or not is_an_image(attributes): + return None + os.makedirs(self.filename, exist_ok=True) + spatial = [int(extent) for extent in shape[1:]] + staging = f"{self.filename}.{name}.{os.getpid()}.tmp.h5" + file, parameters = _create_itk_transform_file(staging, spatial, attributes) + return _ItkTransformDataStream(file, parameters, staging, self._path(name), [3, *spatial]) + + def get_names(self, group: str) -> list[str]: + del group + return sorted({path.stem for pattern in ("*.h5", "*.tfm") for path in Path(self.filename).glob(pattern)}) + + def get_group(self) -> list[str]: + return sorted({path.stem for pattern in ("*.h5", "*.tfm") for path in Path(self.filename).glob(pattern)}) + + def is_exist(self, group: str, name: str | None = None) -> bool: + return os.path.exists(self._path(name if name else group)) + + def get_infos(self, group: str, name: str) -> tuple[list[int], Attribute]: + try: + import h5py + except ImportError: + data, attributes = self.file_to_data(group, name) + return [int(extent) for extent in data.shape], attributes + with h5py.File(self._path(name), "r") as file: + kind = bytes(file["TransformGroup/0/TransformType"][0]) + fixed = np.asarray(file["TransformGroup/0/TransformFixedParameters"][()], dtype=np.float64) + if not kind.startswith(b"DisplacementFieldTransform"): + data, attributes = self.file_to_data(group, name) + return [int(extent) for extent in data.shape], attributes + attributes = Attribute() + attributes["Origin"] = fixed[3:6] + attributes["Spacing"] = fixed[6:9] + attributes["Direction"] = fixed[9:18] + attributes[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" + size_xyz = [int(extent) for extent in fixed[0:3]] + return [3, *size_xyz[::-1]], attributes + class File: def __init__( self, @@ -1893,6 +2106,8 @@ def __enter__(self) -> Dataset.AbstractFile: ) elif self.file_format == "dicom": self.file = Dataset.DicomFile(self.filename, self.read) + elif self.file_format == "itktransform": + self.file = Dataset.ItkTransformFile(self.filename + "/", self.read) else: self.file = Dataset.SitkFile(self.filename + "/", self.read, self.file_format) self.file.__enter__() @@ -2058,6 +2273,8 @@ def can_stream_data(self, attributes: Attribute) -> bool: """ if self.file_format in ("h5", "omezarr"): return True + if self.file_format == "itktransform": + return is_an_image(attributes) return self.file_format == "mha" and is_an_image(attributes) def open_data_stream( diff --git a/tests/unit/test_itk_transform_backend.py b/tests/unit/test_itk_transform_backend.py new file mode 100644 index 00000000..096cd84d --- /dev/null +++ b/tests/unit/test_itk_transform_backend.py @@ -0,0 +1,124 @@ +# 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 + +"""The ``:itktransform`` backend: a displacement field written as an ITK transform file, by regions. + +``sitk.WriteTransform`` needs the whole field resident in float64; the FILE is three HDF5 datasets +that write by regions. Both write paths must be the same file to ITK's reader — same type, fixed +parameters and parameters, exactly — and an entry must read back through ``Dataset.read_transform`` +as the transform it stores.""" + +from pathlib import Path + +import numpy as np +import pytest + +sitk = pytest.importorskip("SimpleITK") +pytest.importorskip("h5py") + +from konfai.utils.dataset import Attribute, Dataset # noqa: E402 + +_ORIGIN, _SPACING = [7.0, -3.0, 10.0], [1.5, 1.5, 2.0] +_DIRECTION = [0.0, -1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0] + + +def _attributes() -> Attribute: + attributes = Attribute() + attributes["Origin"] = np.asarray(_ORIGIN) + attributes["Spacing"] = np.asarray(_SPACING) + attributes["Direction"] = np.asarray(_DIRECTION) + return attributes + + +def _field(seed: int = 0) -> np.ndarray: + return (np.random.default_rng(seed).normal(size=(3, 4, 5, 6)) * 8).astype(np.float32) + + +def _oracle(field: np.ndarray) -> "sitk.DisplacementFieldTransform": + """What ``sitk.WriteTransform`` would have written, held in memory instead.""" + image = sitk.GetImageFromArray(np.moveaxis(field, 0, -1).astype(np.float64), isVector=True) + image.SetOrigin(_ORIGIN) + image.SetSpacing(_SPACING) + image.SetDirection(_DIRECTION) + return sitk.DisplacementFieldTransform(sitk.Cast(image, sitk.sitkVectorFloat64)) + + +def test_the_whole_write_is_the_file_sitk_would_have_written(tmp_path: Path) -> None: + field = _field() + Dataset(tmp_path / "out", "itktransform").write("Transform", "P000", field, _attributes()) + + got = sitk.ReadTransform(str(tmp_path / "out" / "P000" / "Transform.h5")) + want = _oracle(field) + assert got.GetFixedParameters() == want.GetFixedParameters() + assert got.GetParameters() == want.GetParameters() + + +def test_the_streamed_write_is_the_same_file(tmp_path: Path) -> None: + """Region by region, without the field ever whole in RAM — and the same bytes of parameters.""" + field = _field(1) + dataset = Dataset(tmp_path / "out", "itktransform") + stream = dataset.open_data_stream("Transform", "P000", [3, 4, 5, 6], np.dtype("float32"), _attributes()) + assert stream is not None + with stream: + for start in range(0, 4, 2): + stream.write_slice( + (slice(0, 3), slice(start, start + 2), slice(0, 5), slice(0, 6)), field[:, start : start + 2] + ) + + got = sitk.ReadTransform(str(tmp_path / "out" / "P000" / "Transform.h5")) + want = _oracle(field) + assert got.GetFixedParameters() == want.GetFixedParameters() + assert got.GetParameters() == want.GetParameters() + + +def test_an_aborted_stream_leaves_no_entry(tmp_path: Path) -> None: + """A reader must never see a half-written transform under the final name.""" + dataset = Dataset(tmp_path / "out", "itktransform") + stream = dataset.open_data_stream("Transform", "P000", [3, 4, 5, 6], np.dtype("float32"), _attributes()) + assert stream is not None + stream.write_slice((slice(0, 3), slice(0, 2), slice(0, 5), slice(0, 6)), _field()[:, 0:2]) + stream.abort() + + assert not (tmp_path / "out" / "P000" / "Transform.h5").exists() + + +def test_an_entry_reads_back_as_the_transform_it_stores(tmp_path: Path) -> None: + """``read_transform`` and ``get_infos`` answer from the file: the field, its grid, its marker.""" + field = _field(2) + dataset = Dataset(tmp_path / "out", "itktransform") + dataset.write("Transform", "P000", field, _attributes()) + + shape, attributes = dataset.get_infos("Transform", "P000") + assert shape == [3, 4, 5, 6] + np.testing.assert_allclose(attributes.get_np_array("Origin"), _ORIGIN) + + back = dataset.read_transform("Transform", "P000") + want = _oracle(field) + for point in ((8.0, -2.0, 11.0), (7.5, -2.5, 12.0)): + assert back.TransformPoint(point) == pytest.approx(want.TransformPoint(point)) + + +def test_a_foreign_affine_file_reads_back_too(tmp_path: Path) -> None: + """The backend serves any ITK transform file on the read side, not only the fields it writes.""" + affine = sitk.AffineTransform(3) + affine.SetTranslation((2.0, -1.0, 3.0)) + case = tmp_path / "out" / "P000" + case.mkdir(parents=True) + sitk.WriteTransform(affine, str(case / "Reg.h5")) + + back = Dataset(tmp_path / "out", "itktransform").read_transform("Reg", "P000") + point = (1.0, 2.0, 3.0) + assert back.TransformPoint(point) == pytest.approx(affine.TransformPoint(point)) From 05cfac1699400fe6222b5a74205e88b053526329 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 03:09:49 +0200 Subject: [PATCH 27/39] refactor(impact-reg): Transform.h5 is a Write like any other The register run gains a second chain, DVF -> Transform: a plain Write to the ':itktransform' backend, which fills the file region by region. One plan covers the moved image and the transform, resume covers both, and the orchestrator's last hand-rolled writer is gone. Staged .h5 and .tfm transforms route to the same backend, so evaluate's Reg group resolves through Dataset like every other entry. --- .../impact_reg_konfai/impact_reg.py | 99 ++++--------------- .../tests/unit/test_displacement_field_io.py | 19 ---- 2 files changed, 19 insertions(+), 99 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index b9ff3aed..57c71b27 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -52,9 +52,10 @@ _ENSEMBLE_DIR = "Ensemble" -# Writing needs a format named -- nothing is on disk yet to detect one from. Only the store spellings -# need translating; every other suffix is already the token Dataset normalises (".mha" -> "mha"). -_FORMATS = {".ome.zarr": "omezarr", ".zarr": "omezarr"} +# Writing needs a format named -- nothing is on disk yet to detect one from. The store spellings and +# the ITK transform files need translating; every other suffix is already the token Dataset +# normalises (".mha" -> "mha"). +_FORMATS = {".ome.zarr": "omezarr", ".zarr": "omezarr", ".h5": "itktransform", ".tfm": "itktransform"} def _app_id(preset: str) -> str: @@ -216,66 +217,6 @@ def _the_output(dest_dir: Path, stem: str) -> Path: return matches[0] -def _write_displacement_transform(dvf: Path, dest: Path, work: Path) -> None: - """``dest`` as ITK's HDF5 transform writer lays it out — parameters filled region by region. - - ``sitk.WriteTransform`` needs the whole field resident, in float64; the FORMAT does not: an - ITK transform file is three HDF5 datasets (the type; the fixed parameters — size, origin, - spacing, direction; the parameters — the field buffer, component fastest), and HDF5 writes by - regions. The field is read in slabs through ``Dataset``, so the peak is one slab in float64 - instead of the field twice over. Read-back equality with sitk's own writer is pinned by the - app tests. Without ``h5py`` the sitk path serves, whole. - """ - try: - import h5py - except ImportError: - sitk.WriteTransform(sitk.DisplacementFieldTransform(read_displacement_field(dvf)), str(dest)) - return - from konfai.utils.dataset import Dataset - - spec = _stage_group(work / "transform_h5", "DVF", {"P000": dvf}) - filename, _colon, file_format = spec.rpartition(":") - dataset = Dataset(Path(filename), file_format) - shape, attribute = dataset.get_infos("DVF", "P000") - channels, spatial = int(shape[0]), [int(extent) for extent in shape[1:]] - if channels != 3 or len(spatial) != 3: - raise RuntimeError(f"Transform.h5 needs a 3-component 3-D field; '{dvf}' has shape {list(shape)}.") - fixed = np.concatenate( - [ - np.asarray(spatial[::-1], dtype=np.float64), # size, in (x, y, z) - attribute.get_np_array("Origin").astype(np.float64), - attribute.get_np_array("Spacing").astype(np.float64), - attribute.get_np_array("Direction").astype(np.float64).reshape(-1), - ] - ) - rows, offset = 16, 0 - streamed = dataset.bounded_region_reads("DVF", "P000") - resident = None if streamed else np.asarray(dataset.read_data("DVF", "P000")[0]) - staging = dest.with_name(dest.name + ".tmp") - with h5py.File(staging, "w") as file: - file.create_dataset( - "TransformGroup/0/TransformType", - data=[b"DisplacementFieldTransform_double_3_3"], - dtype=h5py.string_dtype(encoding="ascii"), # ITK reads a variable-length ASCII string - ) - file.create_dataset("TransformGroup/0/TransformFixedParameters", data=fixed) - parameters = file.create_dataset( - "TransformGroup/0/TransformParameters", shape=(3 * int(np.prod(spatial)),), dtype=np.float64 - ) - for start in range(0, spatial[0], rows): - stop = min(start + rows, spatial[0]) - if resident is None: - block, _ = dataset.read_data_slice( - "DVF", "P000", (slice(None), slice(start, stop), slice(None), slice(None)) - ) - else: - block = resident[:, start:stop] - slab = np.moveaxis(np.asarray(block, dtype=np.float64), 0, -1).ravel() - parameters[offset : offset + slab.size] = slab - offset += slab.size - os.replace(staging, dest) - - def _run_transform( name: str, datasets: list[str], @@ -502,16 +443,12 @@ def register( self._ensemble_mean(case, presets, dvf_paths, output, work, gpu, cpu, quiet) if not fields_only: - # Every moved image in ONE streamed run: Resample adopts, per case, the grid of that - # case's own DVF (a field is defined ON the fixed grid) and reads the field as the - # map -- one interpolation, slab by slab, the whole cohort under one plan. + # The moved images AND Transform.h5 in ONE streamed run: Resample adopts, per case, + # the grid of that case's own DVF (a field is defined ON the fixed grid) and reads + # the field as the map; the second chain writes the same field as an ITK transform + # file (the ':itktransform' backend fills it region by region) -- one plan, both + # deliverables, resumable. self._derive_moved(dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet) - for case in cases: - # Transform.h5 (consumed by `evaluate` and SlicerImpactReg): the fixed-grid - # field as an ITK transform file, filled region by region from the DVF. - _write_displacement_transform( - _the_output(output / case, "DVF"), output / case / "Transform.h5", work - ) finally: shutil.rmtree(work, ignore_errors=True) @@ -560,13 +497,14 @@ def _derive_moved( cpu: int | None, quiet: bool, ) -> None: - """The moved images, resampled from each moving through ITS displacement field — one run. - - A preset that emits only a field is complete: the moved image IS that field applied to the - moving, so deriving it belongs to this layer. THE GRID AND THE FORMAT FOLLOW THE FIELD, not - the moving: a displacement field is defined ON the fixed grid, so ``reference: '{case}'`` - adopts each case's own DVF grid, and the field is the map (``field_group``) — one - interpolation, streamed, each slab's source window sized from the field values it reads. + """The moved images and ``Transform.h5``, both derived from the fields — one run, two chains. + + A preset that emits only a field is complete: everything else IS that field. The moved + image: ``reference: '{case}'`` adopts each case's own DVF grid (a field is defined ON the + fixed grid) and the field is the map (``field_group``) — one interpolation, streamed, each + slab's source window sized from the field values it reads. ``Transform.h5``: the same field + written as an ITK transform file, a plain ``Write`` the ``:itktransform`` backend fills + region by region. """ from konfai.data.transform import Resample, Write @@ -586,7 +524,8 @@ def _derive_moved( Resample(reference="{case}", reference_group="DVF", field_group="DVF"), Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), ] - } + }, + "DVF": {"Transform": [Write(dataset=f"{output}:itktransform")]}, }, work, gpu, diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index 1194af62..d8efdb30 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -258,22 +258,3 @@ def test_transform_reads_back_identically_from_either_form(tmp_path: Path, suffi reference = sitk.DisplacementFieldTransform(sitk.Image(original)) for point in ((9.0, -1.0, 12.0), (7.5, -2.5, 11.0)): assert restored.TransformPoint(point) == pytest.approx(reference.TransformPoint(point)) - - -@pytest.mark.parametrize("suffix", [".mha", ".ome.zarr"]) -def test_transform_h5_reads_back_as_sitk_would_have_written_it(tmp_path: Path, suffix: str) -> None: - """The streamed writer and ``sitk.WriteTransform`` must be the same file to ITK's reader: - same type, same fixed parameters (size, origin, spacing, direction), same parameters — exactly.""" - from impact_reg_konfai.impact_reg import _write_displacement_transform - - original = _field() - _write_displacement_field(original, tmp_path / f"DVF{suffix}") - work = tmp_path / "work" - work.mkdir() - - _write_displacement_transform(tmp_path / f"DVF{suffix}", tmp_path / "Transform.h5", work) - - got = sitk.ReadTransform(str(tmp_path / "Transform.h5")) - want = sitk.DisplacementFieldTransform(read_displacement_field(tmp_path / f"DVF{suffix}")) - assert got.GetFixedParameters() == want.GetFixedParameters() - assert got.GetParameters() == want.GetParameters() From d4d3083cb3e4fa77b7b8e861cd8e6a1d5dd58829 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 03:43:48 +0200 Subject: [PATCH 28/39] refactor(impact-reg): a registration engine returns its field, not an image Each engine resampled the moving onto the fixed grid and returned it alongside the field, the two concatenated on the channel axis for a pair of ChannelSelect modules to split again. Since a preset declares only its DisplacementField, konfai writes none of it -- but all three engines still computed it, a full-grid sitk.Resample per case, thrown away. The orchestrator derives the moved image from the field anyway, streamed, so the common path resampled twice: once in the engine for nothing, once in impact-reg for real. elastix, fireants and convexadam now return the field alone, the modules stop concatenating, and MovedImage is gone from the three models. --- .../impact_reg_konfai/models/convexadam.py | 21 +++++++--------- .../impact_reg_konfai/models/elastix.py | 10 +++----- .../models/elastix_engine.py | 14 +++++------ .../impact_reg_konfai/models/fireants.py | 25 ++++++++----------- 4 files changed, 30 insertions(+), 40 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/models/convexadam.py b/apps/impact_reg/impact_reg_konfai/models/convexadam.py index c41444ee..7fcbba26 100644 --- a/apps/impact_reg/impact_reg_konfai/models/convexadam.py +++ b/apps/impact_reg/impact_reg_konfai/models/convexadam.py @@ -17,7 +17,7 @@ """ConvexAdam (itk-impact) registration as a self-contained KonfAI model. Same idiomatic ``add_module`` graph and the same output contract as the elastix preset -(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), +(``DisplacementField`` on the FIXED grid), so the orchestrator / app.json / ensemble / uncertainty are unchanged. The engine here is the native, in-memory itk-impact ConvexAdam pipeline (``pip install itk-impact``) instead of the elastix binary: @@ -202,7 +202,7 @@ def _itk_affine_to_sitk(affine: "itk.AffineTransform") -> sitk.AffineTransform: class ConvexAdamEngine: - """Register a fixed/moving pair with the itk-impact ConvexAdam pipeline; return (moved, dvf) on the fixed grid. + """Register a fixed/moving pair with the itk-impact ConvexAdam pipeline; return the displacement field on the fixed grid. The IMPACT feature models are downloaded once (``repo:filename`` on Hugging Face) and reused across cases. Masks are accepted for signature compatibility with the elastix engine but ignored: the ConvexAdam @@ -427,8 +427,8 @@ def register( device_index: int, fixed_mask: sitk.Image | None = None, moving_mask: sitk.Image | None = None, - ) -> tuple[np.ndarray, np.ndarray]: - """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.""" + ) -> np.ndarray: + """Register ``moving`` onto ``fixed``; return the displacement field, channel-first, on the fixed grid.""" device = f"cuda:{device_index}" if device_index >= 0 else "cpu" fixed_itk = _sitk_to_itk(fixed) moving_itk = _sitk_to_itk(moving) @@ -452,7 +452,6 @@ def register( if field is not None: chain.append(_itk_field_to_sitk_transform(field, fixed)) composite = sitk.CompositeTransform(chain) - moved = sitk.Resample(moving, fixed, composite, sitk.sitkLinear, 0.0, moving.GetPixelID()) dvf = sitk.TransformToDisplacementField( composite, sitk.sitkVectorFloat64, @@ -461,9 +460,8 @@ def register( fixed.GetSpacing(), fixed.GetDirection(), ) - moved_np, _ = image_to_data(moved) dvf_np, _ = image_to_data(dvf) - return moved_np, dvf_np + return dvf_np class ConvexAdamRegistration(torch.nn.Module): @@ -500,8 +498,8 @@ def forward( for b in range(fixed.shape[0]): fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b]) moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) - moved_np, dvf_np = self._engine.register(fixed_img, moving_img, device_index) - combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + dvf_np = self._engine.register(fixed_img, moving_img, device_index) + combined.append(torch.from_numpy(dvf_np)) return torch.stack(combined, dim=0).to(fixed.device) @@ -521,7 +519,7 @@ class RegistrationNet(network.Network): """Pairwise ConvexAdam registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1; the mask branches 2/3 are accepted but unused by this engine). - Outputs on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` (the + Output on the fixed grid: ``DisplacementField`` (the DIM-component displacement field, in mm). Geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. """ @@ -629,5 +627,4 @@ def __init__( self.add_module( "Registration", ConvexAdamRegistration(engine), in_branch=[0, 1, 2, 3], out_branch=["registration"] ) - self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) - self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) + self.add_module("DisplacementField", ChannelSelect(0, 3), in_branch=["registration"], out_branch=["dvf"]) diff --git a/apps/impact_reg/impact_reg_konfai/models/elastix.py b/apps/impact_reg/impact_reg_konfai/models/elastix.py index 3c5603c6..c2d72b0f 100644 --- a/apps/impact_reg/impact_reg_konfai/models/elastix.py +++ b/apps/impact_reg/impact_reg_konfai/models/elastix.py @@ -17,7 +17,7 @@ """Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph. ``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks = -2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns +2/3) and emits its ``DisplacementField`` on the fixed grid. This module owns the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download, subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built. @@ -279,9 +279,8 @@ class RegistrationNet(network.Network): """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2, moving mask = 3; masks restrict the metric, whole-image = no restriction). - Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` - (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two - ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via + Output, on the fixed grid: ``DisplacementField`` + (the dim-component displacement field, mm). Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. """ @@ -366,5 +365,4 @@ def __init__( in_branch=[0, 1, 2, 3], out_branch=["registration"], ) - self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) - self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) + self.add_module("DisplacementField", ChannelSelect(0, 3), in_branch=["registration"], out_branch=["dvf"]) diff --git a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py index 60018640..91b2a1ad 100644 --- a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py +++ b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py @@ -56,7 +56,7 @@ def _is_partial_mask(mask: "sitk.Image | None") -> bool: class ElastixEngine: - """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid. + """Run the elastix-IMPACT binary on a fixed/moving pair; return the displacement field on the fixed grid. NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``. @@ -253,8 +253,8 @@ def register( device_index: int, fixed_mask: sitk.Image | None = None, moving_mask: sitk.Image | None = None, - ) -> tuple[np.ndarray, np.ndarray]: - """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid. + ) -> np.ndarray: + """Register ``moving`` onto ``fixed``; return the displacement field, channel-first, on the fixed grid. Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none. @@ -353,7 +353,6 @@ def register( raise FileNotFoundError("elastix produced no composite transform file.") transform = sitk.ReadTransform(str(transforms[-1])) - moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) dvf = sitk.TransformToDisplacementField( transform, sitk.sitkVectorFloat64, @@ -362,9 +361,8 @@ def register( fixed.GetSpacing(), fixed.GetDirection(), ) - moved_np, _ = image_to_data(moved) dvf_np, _ = image_to_data(dvf) - return moved_np, dvf_np + return dvf_np finally: shutil.rmtree(work, ignore_errors=True) @@ -424,8 +422,8 @@ def forward( moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) - moved_np, dvf_np = self._engine.register( + dvf_np = self._engine.register( fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img ) - combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + combined.append(torch.from_numpy(dvf_np)) return torch.stack(combined, dim=0).to(fixed.device) diff --git a/apps/impact_reg/impact_reg_konfai/models/fireants.py b/apps/impact_reg/impact_reg_konfai/models/fireants.py index 6220e7ec..853fcac5 100644 --- a/apps/impact_reg/impact_reg_konfai/models/fireants.py +++ b/apps/impact_reg/impact_reg_konfai/models/fireants.py @@ -22,7 +22,7 @@ """FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets). Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset -(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the +(``DisplacementField`` on the FIXED grid), so the orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages: @@ -53,7 +53,7 @@ The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly. -``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK — +the emitted ``DisplacementField`` is rebuilt from that transform with SimpleITK — the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here). @@ -437,7 +437,7 @@ def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor: class FireANTsEngine: """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return - (moved, dvf) on the fixed grid. + the displacement field on the fixed grid. ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet`` @@ -583,8 +583,8 @@ def register( device_index: int, fixed_mask: sitk.Image | None = None, moving_mask: sitk.Image | None = None, - ) -> tuple[np.ndarray, np.ndarray]: - """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.""" + ) -> np.ndarray: + """Register ``moving`` onto ``fixed``; return the displacement field, channel-first, on the fixed grid.""" ensure_fireants_runtime() from fireants.io import BatchedImages, Image from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones @@ -714,9 +714,8 @@ def register( if torch.cuda.is_available(): torch.cuda.synchronize() - # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path, + # The DVF is rebuilt from the single transform on the fixed grid — the ConvexAdam output path, # so every FireANTs preset emits identical-shaped results. - moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID()) dvf = sitk.TransformToDisplacementField( transform, sitk.sitkVectorFloat64, @@ -725,9 +724,8 @@ def register( fixed.GetSpacing(), fixed.GetDirection(), ) - moved_np, _ = image_to_data(moved) dvf_np, _ = image_to_data(dvf) - return moved_np, dvf_np + return dvf_np class FireANTsRegistration(torch.nn.Module): @@ -768,10 +766,10 @@ def forward( moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) - moved_np, dvf_np = self._engine.register( + dvf_np = self._engine.register( fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img ) - combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0))) + combined.append(torch.from_numpy(dvf_np)) return torch.stack(combined, dim=0).to(fixed.device) @@ -791,7 +789,7 @@ class RegistrationNet(network.Network): """Pairwise FireANTs registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2, moving mask = 3; masks restrict the metric, whole-image = no restriction). - Outputs on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField`` + Output on the fixed grid: ``DisplacementField`` (the DIM-component displacement field, in mm). Geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``. The knobs below are read straight from these annotations by the UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method`` @@ -908,5 +906,4 @@ def __init__( self.add_module( "Registration", FireANTsRegistration(engine), in_branch=[0, 1, 2, 3], out_branch=["registration"] ) - self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"]) - self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"]) + self.add_module("DisplacementField", ChannelSelect(0, 3), in_branch=["registration"], out_branch=["dvf"]) From be3bdda03b5ae17fae3cd6613814366df06ae099 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 04:27:03 +0200 Subject: [PATCH 29/39] feat(data)!: drop the recorded field bound MaxDisplacement is gone: the OME-Zarr writer records nothing, the streamed field write accumulates nothing, and the reader scans nothing -- the run sizes every region's pull from the field values it reads for sampling, and the plan prices those reads as a zero field and says so in its note. What remains of the plan-time scan is the one thing it still owes: every field entry's header must open, or the group answers whole-volume before a case is chosen. BREAKING CHANGE: DISPLACEMENT_BOUND_ATTRIBUTE and displacement_bound are removed from konfai.utils.ome_zarr; stores carrying the attribute simply have it ignored. --- docs/source/config_guide/transform.md | 20 +- .../source/reference/components/transforms.md | 2 +- konfai/data/transform.py | 152 +++++------ konfai/utils/dataset.py | 30 +-- konfai/utils/ome_zarr.py | 30 --- tests/unit/test_dataset.py | 6 +- tests/unit/test_resample_to_reference.py | 45 +--- tests/unit/test_warp.py | 72 ++---- tests/unit/test_write_pyramid.py | 133 ++++++++++ .../test_write_pyramid_and_field_bound.py | 241 ------------------ 10 files changed, 224 insertions(+), 507 deletions(-) create mode 100644 tests/unit/test_write_pyramid.py delete mode 100644 tests/unit/test_write_pyramid_and_field_bound.py diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index 76c9aeba..6fc9c690 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -103,7 +103,7 @@ few percent of the budget can still exceed it. | --- | --- | | `allow` | Take the whole-volume path silently — but the plan still names it. | | `warn` (default) | Same, plus a warning line after the plan. | -| `error` | Refuse the run. Nothing is written. A fallback only discovered mid-run (a failed sweep, a `Resample` field bound exceeded) stops at that case — earlier cases stay written, and the per-case resume covers the rerun. | +| `error` | Refuse the run. Nothing is written. A fallback only discovered mid-run (a failed sweep) stops at that case — earlier cases stay written, and the per-case resume covers the rerun. | Independently of `on_fallback`, a case that **cannot stream and does not fit `memory_budget`** always refuses the whole run, before the first byte. Writing @@ -365,17 +365,13 @@ 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. -Nothing is declared about how far the field reaches. The field window a -region samples is its own box, read for sampling regardless — and the sup of -the values just read bounds that region's source pull, so each slab pays -exactly the halo *its* displacements require, measured at run from a read the -sampler needed anyway. A bound the **store recorded at write time** (KonfAI -records one on the OME-Zarr fields it writes) does two things without anyone -asking: it lets the plan **price** the reads exactly — without one the -estimate assumes a zero field, and the plan says so — and it is **checked -against every field region actually read**: a store whose data exceed its own -metadata raises rather than sampling zeros, which would show up as a dark rim -around the moved anatomy and nothing else. +Nothing is declared about how far the field reaches, and nothing is recorded. +The field window a region samples is its own box, read for sampling regardless +— and the sup of the values just read bounds that region's source pull, so +each slab pays exactly the halo *its* displacements require, measured at run +from a read the sampler needed anyway. The one thing the plan cannot know from +headers is the cost of those reads: it prices them as if the field were zero, +and says so. Naming no target grid is the shape update of an atlas build — the field applied on the case's *own* grid — and is the same stage with `reference` left out: diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 03b87897..83517586 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -97,7 +97,7 @@ until it declares otherwise. | --- | --- | --- | --- | --- | --- | | `Padding` | `F.pad`; updates Origin. `mode` supports `"constant:"`. | `padding=[0,0,0,0,0,0], mode="constant", inverse=True` | **yes** | **yes** | no‡ | | `Crop` | Crop to foreground bounding box; caches the box; updates Origin. | `inverse=True` | **yes** | **yes** (pads back) | **yes** — once the `box` is on the case; the region is the patch translated | -| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region — and a bound the store recorded at write time prices the plan exactly and is **checked** per component against every region read. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | +| `Resample` | **The one resample.** Two questions: which grid to write on — nothing (the case's own), `spacing`, `shape`, or `reference` (a stored image's grid, adopted whole) — and what map to write it through — `field` (a displacement field on its own grid, in world units) and/or `transforms` (rigid, affine, BSpline, dense field or composite stored beside the cases; the **last declared is applied first**). Asked for together they compose into **one interpolation**. `align` places a `spacing`/`shape` grid: `extent` keeps the field of view (the outer faces coincide), `origin` keeps voxel zero's centre. `interpolation` left unset is nearest for `uint8` and linear otherwise. | `spacing=None`, `shape=None`, `reference=None`, `reference_group=None`, `reference_dataset=None`, `transforms=None`, `field=None`, `field_group=None`, `align="extent"`, `interpolation=None`, `fill=0.0`, `inverse=True` | **yes** | **yes** — the grid change alone; a declared map is not inverted, and a stage that changes no grid refuses rather than pretend | **yes** — declares `REGRID`. A rigid or affine map bounds exactly; a BSpline and a stored field bound by the sup-norm of their values, which is a theorem (non-negative kernels summing to one), not a sample of the boundary; a field on disk sizes each region's pull from the field values read for sampling anyway — measured at run, per region; the plan prices those reads as a zero field, and says so. Falls back with the reason when the case carries no geometry, a type decomposes into no bounded map, or `invert: true` names a spline or a field | | `Canonical` | Reorient to canonical direction (3-D); updates Origin/Direction. | `inverse=True` | **yes** — a remap that transposes extents moves the patch grid | **yes** | **yes** — when the case's direction is a signed axis permutation; no on an oblique one (it is resampled) | | `Permute` | Permute spatial axes. `dims` is a pipe-separated axis list. | `dims="1\|0\|2", inverse=True` | **yes** | **yes** | **yes** — index remap | | `Flip` | Flip spatial axes. | `dims="1\|0\|2", inverse=True` | no | **yes** (self-inverse) | **yes** — index remap | diff --git a/konfai/data/transform.py b/konfai/data/transform.py index ad7a2655..39adf447 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -1204,9 +1204,8 @@ class Resample(TransformInverse): -- a theorem, not a sample of the boundary. A field on disk is read region by region, and the window a region samples is its own box: the sup of the values just read bounds that region's pull, so each slab pays exactly the halo ITS displacements require -- measured at run, from a - read the sampler needs regardless. Nothing is declared: a bound the STORE recorded at write - time (KonfAI's OME-Zarr fields carry one) prices the plan exactly and is CHECKED against every - region read; without one the plan prices the reads as if the field were zero, and says so. + read the sampler needs regardless. Nothing is declared and nothing is recorded: the plan + prices the reads as if the field were zero, and says so. ``align`` decides where a ``spacing`` or a ``shape`` grid SITS, and it is the one silent choice in the family -- a quarter of a voxel of anatomy, made differently by every library that offers @@ -1458,7 +1457,6 @@ def _field_stage(self, name: str, region: Grid) -> DisplacementStage: 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)) - source.check_bound(values, name) stage = DisplacementStage(grid.sub_grid(window), values.numpy(), order=1) self._field_window = (name, key, stage) return stage @@ -1477,17 +1475,13 @@ def _stages(self, name: str, region: Grid) -> SpatialStages: return tuple(stages) def _bound(self, name: str) -> TransformBound: - """What the map is guaranteed to do — from recorded bounds and coefficients, no voxel read.""" + """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: - recorded = self.displacement.component_bound() - if recorded is None: - raise TransformError( - "the field carries no recorded bound, so what it is guaranteed to do is unknown" - " before its values are read." - ) - folded = TransformBound.shift(np.asarray(recorded[:rank], dtype=np.float64)).after(folded) + raise TransformError( + "a field's reach is unknown before its values are read; nothing bounds it from headers." + ) if self.transforms is not None: folded = bound_of(self._stored_stages(name), rank).after(folded) return folded @@ -1495,12 +1489,11 @@ 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 with no declared or recorded bound prices as zero displacement. The run never - trusts this window: a declared 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. 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. """ - if self.displacement is None or self.displacement.component_bound() is not None: + if self.displacement is None: return self._bound(name) rank = self._source_grid(name).rank folded = TransformBound.exact(AffineMap.identity(rank)) @@ -1529,6 +1522,8 @@ def _require_runnable(self, name: str) -> None: bytes are written. ``transform_shape`` runs for every case as the plan is built, which is the earliest the failure is knowable and the only place it costs nothing. """ + if self.displacement is not None: + self.displacement.probe(name) if self.transforms is None: return try: @@ -1583,8 +1578,7 @@ def _probe_cohort(self) -> str | None: # under it fails both routes on whichever case reaches it. A field that merely records no # bound streams: its windows are sized from the values the run reads (measured_region_source). if self.displacement is not None: - self.displacement.component_bound() - if self.displacement.scan_failed: + if not self.displacement.headers_readable(): return ( "an entry in the field group could not be header-read, so what any region of it" " must pull is unknown. Check the field store: one unreadable entry anywhere" @@ -1614,14 +1608,8 @@ def stream_region_source( @property def measures_at_run(self) -> bool: - """Whether the run sizes this stage's windows from the data it reads. - - Only a field with NO declared or recorded bound: a bounded field keeps the declared - window, whose streamed result is bit-identical to the whole-volume path on a separable - map — measuring would tighten its windows at the price of that identity. Measuring is the - route for the field that could not stream at all before. - """ - return self.displacement is not None and self.displacement.component_bound() is None + """Whether the run sizes this stage's windows from the data it reads — any declared field.""" + return self.displacement is not None def measured_region_source( self, name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute @@ -1631,8 +1619,7 @@ def measured_region_source( The field window a region needs is its own box, read for sampling regardless; the sup of the values just read bounds every interpolated displacement in the region (a convex combination cannot exceed the lattice values it blends), so the window is exact per region - — a quiet slab pays a quiet halo. A bound the store recorded stays the cap those values - are checked against. + — a quiet slab pays a quiet halo. """ del source_spatial_shape, cache_attribute source, target = self._grids_of(name) @@ -1738,11 +1725,13 @@ def _map_bound(self, name: str) -> TransformBound | None: ``None`` when there is no map — or when nothing bounds it: a coverage that cannot be judged must not refuse, and the unboundable configurations carry a fallback reason of their own. + A field prices as zero displacement here (:meth:`_pricing_bound`), so a stored affine + beside it still places the samples instead of the grids being judged bare. """ if self.transforms is None and self.displacement is None: return None try: - return self._bound(name) + return self._pricing_bound(name) except Exception: # an unreadable or unbounded map answers None, never a crash return None @@ -1784,8 +1773,11 @@ def _refuse_if_disjoint(self, name: str) -> None: 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. + + 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. """ - if self._target_is_own or self.coverage(name) > 0.0: + if self._target_is_own or self.displacement is not None or self.coverage(name) > 0.0: return where = f"case '{name}'" if name else "the case" raise TransformError( @@ -1806,13 +1798,12 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut """ del group_dest notes: list[str] = [] - if self.displacement is not None and self.displacement.component_bound() is None: + if self.displacement is not None: # Case-independent on purpose: the plan prints identical notes once, so this is one line # for the stage rather than one per case. notes.append( - "the field carries no bound, so each region's source window is sized from the field" - " values read at run; the read estimate prices the field as zero. A field with a" - " recorded bound (KonfAI records one on the OME-Zarr fields it writes) prices exactly" + "each region's source window is sized from the field values read at run; the read" + " estimate prices the field as zero" ) try: source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") @@ -2305,47 +2296,29 @@ def __init__(self, field: str | None, group: str | None) -> None: self.group = group #: The run's own roots, handed over by the owner; only consulted when there is no path. self.roots: list[Dataset] = [] - #: Whether the header scan DIED, as opposed to finding no recorded bound: an unreadable - #: entry fails both routes at run, where a merely bound-less field streams (measured). - self.scan_failed = False - self._recorded_bound: list[float] | None = None - self._scan_resolved = False - - def component_bound(self) -> list[float] | None: - """The per-component bound the STORE recorded, or ``None`` when it recorded none. - - The largest bound any field in the group recorded at write time, read from headers alone - and memoized — nobody declares anything. If a single entry carries no bound the answer is - ``None``: a maximum over the others would be a bound for them and a guess for that one. - Per component, in the field's own (x, y, z) order, because these grids are anisotropic: - one collapsed maximum over-reads the fine axes. - """ - if self._scan_resolved: - return self._recorded_bound - self._scan_resolved = True - from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE + self._scan_ok: bool | None = None + self._probed: set[str] = set() - bound: list[float] = [] + def headers_readable(self) -> bool: + """Whether every field entry's HEADER opens, memoized — the plan's one probe of the group. + + An unreadable entry fails both routes on whichever case reaches it, and a directory store + lists its entries from the filesystem alone, so a corrupt one only surfaces at its header: + scanned here, one entry at a time, before any case is chosen. + """ + if self._scan_ok is not None: + return self._scan_ok try: group = self.group_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 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._recorded_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)] + root.get_infos(group, entry) except Exception: # an unreadable field dataset is a whole-volume answer, not a crash - self.scan_failed = True - return self._recorded_bound - if bound and max(bound) > 0.0: - self._recorded_bound = bound - return self._recorded_bound + self._scan_ok = False + return False + self._scan_ok = True + return True def group_for(self, name: str | None) -> str: if self.group is not None: @@ -2383,6 +2356,26 @@ 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 probe(self, name: str) -> None: + """The case's own entry, proven present and readable — the per-case half of the scan. + + :meth:`headers_readable` walks what is on disk; it cannot know which cases the plan will + ask for, so a missing entry would otherwise surface mid-run, after bytes are written. + Memoized: one header read per case, at plan time. + """ + if name in self._probed: + return + try: + self.infos(name) + except TransformError: + raise + except Exception as error: + raise TransformError( + f"'Resample' cannot read the field header for case '{name}': {type(error).__name__}: {error}.", + "Repair or re-write that entry, or drop the case with 'subset'.", + ) from error + self._probed.add(name) + def read(self, name: str, region: tuple[slice, ...] | None, channels: int) -> torch.Tensor: group = self.group_for(name) root = self._root_for(name) @@ -2399,27 +2392,6 @@ def read(self, name: str, region: tuple[slice, ...] | None, channels: int) -> to ) return field - def check_bound(self, field: torch.Tensor, name: str) -> None: - """The store's recorded bound is a promise about the region 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]): - recorded = bound[component] if component < len(bound) else max(bound) - largest = float(field[component].abs().max()) - if largest > recorded: - raise TransformError( - f"The field for case '{name}' displaces up to {largest:.3f} on component" - f" {component}, beyond the {recorded:.3f} its store recorded — the bound" - " 'Resample' sized its region from.", - "The store's metadata contradicts its data: rewrite the field so the recorded" - " bound holds, or strip the stale bound so the windows are measured instead.", - ) - class Reduce(Transform): """Fold every case of a group into one volume, at fixed voxel. diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 97a2c39e..ad4fffc3 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -861,46 +861,25 @@ def __init__( final_path: Path, scale_factors: list[int] | None = None, downsample_method: str | None = None, - displacement_field: bool = False, ) -> None: self._array = array self._store_path = store_path self._final_path = final_path self._scale_factors = scale_factors self._downsample_method = downsample_method - self._displacement_field = displacement_field - # Running per-component bound of a streamed field. Accumulated from the regions as they are - # written -- the only place the samples are ever all seen, since the point of this path is - # that the field never exists whole. - self._bound: list[float] = [] def write_slice(self, slices: tuple[slice, ...], data: np.ndarray) -> None: self._array[slices] = data - if self._displacement_field: - from konfai.utils.ome_zarr import displacement_bound - - block = displacement_bound(data) - if len(self._bound) < len(block): - self._bound.extend([0.0] * (len(block) - len(self._bound))) - for component, value in enumerate(block): - self._bound[component] = max(self._bound[component], value) def _close(self, success: bool) -> None: from konfai.utils.ome_zarr import ( - DISPLACEMENT_BOUND_ATTRIBUTE, append_ome_zarr_levels, clear_ome_zarr_cache, - update_konfai_attributes, ) if not success: shutil.rmtree(self._store_path, ignore_errors=True) return - if self._displacement_field and self._bound: - # Recorded before the levels are derived and before the rename, so the store is published - # complete: a consumer that finds the entry finds its bound with it, or the entry is not - # there at all. - update_konfai_attributes(self._store_path, {DISPLACEMENT_BOUND_ATTRIBUTE: self._bound}) if self._scale_factors: # On the temporary store, so the rename below publishes level 0 and its coarser levels in # one step. append_ome_zarr_levels REWRITES level 0 (ngff-zarr composes a multiscales as a @@ -1761,14 +1740,7 @@ def open_data_stream( # so the stream derives it at finalize, on the TEMPORARY store, before the rename. That # order is what keeps publication atomic: a reader never sees a store whose level 0 is # complete but whose coarser levels are not. - return _OmeZarrDataStream( - array, - store_path, - final_path, - self.scale_factors, - self.downsample_method, - DISPLACEMENT_FIELD_ATTRIBUTE in attributes, - ) + return _OmeZarrDataStream(array, store_path, final_path, self.scale_factors, self.downsample_method) def get_names(self, group: str) -> list[str]: return self.get_group() diff --git a/konfai/utils/ome_zarr.py b/konfai/utils/ome_zarr.py index 2a9d5c95..1f834c16 100644 --- a/konfai/utils/ome_zarr.py +++ b/konfai/utils/ome_zarr.py @@ -88,10 +88,6 @@ def _native_byteorder(array: np.ndarray) -> np.ndarray: # everywhere else, being the version portable across the whole CI matrix. _DISPLACEMENT_AXIS_TYPE = "displacement" -#: Largest absolute displacement per component, world units, recorded by the producer of a field. -#: A consumer sizes the region it must read from this, and reading it costs a header where measuring -#: it costs a scan of the whole field. -DISPLACEMENT_BOUND_ATTRIBUTE = "MaxDisplacement" _RFC5_VERSION = "0.6" _DEFAULT_VERSION = "0.4" @@ -400,37 +396,11 @@ def write_ome_zarr( ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version) recorded = dict(attributes) if attributes else {} - if displacement_field: - # The producer holds the samples, so the bound is free here and a full scan anywhere else. - recorded[DISPLACEMENT_BOUND_ATTRIBUTE] = displacement_bound(data) if recorded: group = zarr.open_group(str(store_path), mode="r+") group.attrs[_KONFAI_ATTR_KEY] = {"attributes": recorded} -def displacement_bound(data: np.ndarray) -> list[float]: - """The largest absolute displacement per component, in the field's own world units. - - This is the number a consumer needs to size the region it must read before resampling through the - field, and it is a property of the DATA, not of any parameter -- so the only place it is free is - here, where the producer already holds the samples. Read back from the store's attributes it costs - a header; recomputed by the consumer it costs a full scan of a field that can be 13.6 GiB. - - Per component, not one scalar: the component axis is (x, y, z) where the array axes are (z, y, x), - and a consumer turns each bound into voxels by its OWN axis spacing -- these grids are anisotropic - (40 um in z against 30.08 in x/y here), so a single collapsed maximum over-reads two axes. - - Computed in float32, the dtype a field is stored in: a bound rounded up in float64 and then - compared against float32 samples is the one failure mode this is meant to remove. - """ - field = np.asarray(data) - if field.ndim < 2: - return [] - # Cast before the max, not after: ``initial=np.float32(0.0)`` does not downcast a float64 input. - flat = np.abs(field.reshape(field.shape[0], -1)).astype(np.float32, copy=False) - return [float(flat[component].max(initial=np.float32(0.0))) for component in range(flat.shape[0])] - - def update_konfai_attributes(store_path: str | Path, extra: dict[str, Any]) -> None: """Merge ``extra`` into the store's KonfAI attribute sidecar, keeping what is already there. diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py index 64c82144..3b369b28 100644 --- a/tests/unit/test_dataset.py +++ b/tests/unit/test_dataset.py @@ -74,17 +74,17 @@ def test_attribute_repeated_set_returns_latest_version() -> None: def test_attribute_built_from_a_store_sidecar_holds_text_a_writer_accepts() -> None: - """An OME-Zarr sidecar is JSON, so it hands back live lists -- ``MaxDisplacement`` is one. + """An OME-Zarr sidecar is JSON, so it hands back live lists, not their string form. Both doors normalize to text, construction included: a value deep-copied through construction untouched reaches ``Image.SetMetaData``, which accepts only ``std::string`` -- a field that can be written but never reopened. """ - attribute = Attribute({"MaxDisplacement": [1.19, 2.39, 3.59], "Spacing": (1.5, 1.5, 2.0)}) + attribute = Attribute({"WorldReach": [1.19, 2.39, 3.59], "Spacing": (1.5, 1.5, 2.0)}) assert all(isinstance(value, str) for value in attribute.values()) # And readable back: a list prints comma-separated, which np.fromstring alone could not read. - np.testing.assert_allclose(attribute.get_np_array("MaxDisplacement"), [1.19, 2.39, 3.59]) + np.testing.assert_allclose(attribute.get_np_array("WorldReach"), [1.19, 2.39, 3.59]) np.testing.assert_allclose(attribute.get_np_array("Spacing"), [1.5, 1.5, 2.0]) diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 8e3b6c61..26be6065 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -38,7 +38,6 @@ from konfai.data.transform import LocalityKind, Reduce, Resample, 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 @@ -858,30 +857,9 @@ def refuse(*args: object, **kwargs: object) -> None: assert list(written.shape[1:]) == list(_REFERENCE_SPATIAL) -def test_a_field_beyond_its_recorded_bound_is_refused( - warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path -) -> None: - """The store's bound is a promise about what was read; data that break 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 - attributes = _attributes(_FIELD_ORIGIN, _FIELD_SPACING) - attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.asarray([0.5, 0.5, 0.5]) - lying = Dataset(tmp_path / "lying", "mha") - lying.write("DVF", _CASE, _displacement(), attributes) - stage = Resample( - reference=_CASE, reference_group="Reference", field=f"{tmp_path / 'lying'}:mha", field_group="DVF", fill=_FILL - ) - stage.set_datasets([images]) - 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_still_streams(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: """The run sizes each region's pull from the field - values it reads for sampling, so a missing recorded bound is a pricing gap, not a fallback.""" + values it reads for sampling: the plan prices the field as zero -- a pricing gap, not a fallback.""" images, fields, _volume = warped stage = _warping(images, fields) assert stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind is LocalityKind.REGRID @@ -938,27 +916,6 @@ 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_the_recorded_bound_is_read_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 = Resample(reference="CASE_1", reference_group="Reference", field_group="DVF") - 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_the_two_gathers_obey_the_same_rules_through_an_identity_field(tmp_path: Path) -> None: """One arithmetic, two loops: per-axis maps, and eight corners at a coordinate volume. diff --git a/tests/unit/test_warp.py b/tests/unit/test_warp.py index 9bf9edd1..c252eb10 100644 --- a/tests/unit/test_warp.py +++ b/tests/unit/test_warp.py @@ -17,8 +17,7 @@ """``Resample`` through a displacement field alone: a warp on the case's own grid, region by region. The claim under test is the one that matters for a volume larger than memory: the streamed result -equals the whole-volume one, each region's window is sized from the field values it reads, and a -bound the store recorded is verified rather than trusted.""" +equals the whole-volume one, and each region's window is sized from the field values it reads.""" from pathlib import Path @@ -26,10 +25,10 @@ import pytest import torch from konfai.data.patching import DatasetManager -from konfai.data.transform import LocalityKind, RegionContext, Resample, Save -from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, Attribute, Dataset +from konfai.data.transform import LocalityKind, Resample, Save +from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import TransformError -from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE, _zarr_v3_available +from konfai.utils.ome_zarr import _zarr_v3_available pytest.importorskip("SimpleITK") @@ -135,32 +134,6 @@ def test_an_oblique_case_grows_its_window_on_every_axis(tmp_path: Path) -> None: assert all(width > 1 for width in widths), f"a turned case reaches on every axis, got {widths}" -@_needs_rfc5 -def test_the_bound_the_fields_recorded_prices_the_plan(tmp_path: Path) -> None: - """The recorded bound is the number the producer already knew — read from headers, declared by - nobody. It sizes the plan's (headers-only) windows; per component and not one collapsed - maximum, because these grids are anisotropic and one number over-reads the fine axes.""" - fields = Dataset(tmp_path / "dvf", "omezarr") - for case, shift in (("CASE_000", (1.0, 2.0, 3.0)), ("CASE_001", (0.5, 6.0, 1.0))): - field = np.zeros((3, 4, 5, 6), dtype=np.float32) - for component, value in enumerate(shift): - field[component] = value - attribute = _attributes() - attribute[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" - fields.write("DVF", case, field, attribute) - - warp = Resample(field=f"{tmp_path / 'dvf'}:omezarr", field_group="DVF") - locality = warp.patch_locality(_attributes()) - - # The cohort's bound is (x=1.0, y=6.0, z=3.0); spacing in array order (z, y, x) is (1, 1, 2), so - # the window grows by 3, 6 and 1 voxels (plus the linear taps' one) around its target. - assert locality.kind is LocalityKind.REGRID - window = _recorded(warp).stream_region_source("CASE_000", (slice(4, 6),) * 3, [10, 12, 14], _attributes()) - reaches = (3.0, 6.0, 0.5) # array order (z, y, x): the bound divided by that axis's spacing - starts = [max(0, int(np.floor(4 - 0.5 - reach)) - 1) for reach in reaches] - assert [part.start for part in window] == starts - - def test_the_header_scan_survives_an_unreadable_entry_in_the_field_group(tmp_path: Path) -> None: """The whole group is header-read, including entries this run never warps. @@ -249,32 +222,6 @@ def test_streamed_equals_whole_volume(tmp_path: Path, monkeypatch: pytest.Monkey np.testing.assert_allclose(streamed, reference, rtol=1e-5, atol=1e-4) -def test_a_field_beyond_its_recorded_bound_raises(tmp_path: Path) -> None: - """The store's metadata is a promise about what was read; a store whose data break it must not - sample zeros — a dark rim around the moved anatomy and nothing else to see. - - Checked per component, the way the halo is derived: a field under the collapsed maximum can - still exceed the bound on one axis, and that axis is the one whose halo was too small. - """ - _source, _fields, volume = _fixture(tmp_path, shift_um=(0.0, 0.0, 9.0)) - attributes = _attributes() - attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.asarray([1.0, 1.0, 1.0]) - field = np.zeros((3, 10, 12, 14), dtype=np.float32) - field[2] = 9.0 - Dataset(tmp_path / "lying", "mha").write("DVF", "CASE_000", field, attributes) - warp = Resample(field=f"{tmp_path / 'lying'}:mha", field_group="DVF") - - _recorded(warp) - with pytest.raises(TransformError, match=r"on component 2, beyond the 1\.000"): - whole = (slice(0, 10), slice(0, 12), slice(0, 14)) - warp.stream_region( - "CASE_000", - torch.from_numpy(volume), - RegionContext(whole, whole, (10, 12, 14), (10, 12, 14)), - _attributes(), - ) - - def test_a_field_with_the_wrong_component_count_is_named(tmp_path: Path) -> None: rng = np.random.default_rng(1) Dataset(tmp_path / "src", "h5").write("CT", "CASE_000", rng.random((1, 4, 4, 4)).astype(np.float32), _attributes()) @@ -285,6 +232,17 @@ def test_a_field_with_the_wrong_component_count_is_named(tmp_path: Path) -> None warp("CASE_000", torch.zeros(1, 4, 4, 4), _attributes()) +def test_a_case_with_no_field_entry_is_refused_at_plan_time(tmp_path: Path) -> None: + """The cohort scan proves what is on disk; only the per-case probe knows what the plan asks for. + + Without it a missing entry surfaces mid-run, in the field read, after bytes are written. + """ + _fixture(tmp_path) # writes a field for CASE_000, and for no other case + warp = Resample(field=f"{tmp_path / 'dvf'}:h5", field_group="DVF") + with pytest.raises(TransformError, match="CASE_MISSING"): + warp.transform_shape("CT", "CASE_MISSING", [10, 12, 14], _attributes()) + + def test_an_empty_field_path_declares_no_field() -> None: """An empty ``field`` with no group is the identity map, not a broken declaration.""" assert Resample(field="").displacement is None diff --git a/tests/unit/test_write_pyramid.py b/tests/unit/test_write_pyramid.py new file mode 100644 index 00000000..871664cd --- /dev/null +++ b/tests/unit/test_write_pyramid.py @@ -0,0 +1,133 @@ +# 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 + +"""The write side of the data surface: a declared pyramid, written by both paths identically. + +A pyramid is indexed by position (``:omezarr@1``), so a producer that writes one level where two +were promised does not fail -- it resolves ``@1`` to something else, or to nothing. So each test +here asserts the two write paths agree: assembled in memory, and streamed region by region. They +are different code (``write_ome_zarr`` against ``create_ome_zarr_store`` + +``append_ome_zarr_levels``), and the streamed one is the path a real volume takes. +""" + +from pathlib import Path + +import numpy as np +import pytest +import zarr +from konfai.utils.dataset import Attribute, Dataset +from konfai.utils.errors import DatasetManagerError +from konfai.utils.ome_zarr import ( + append_ome_zarr_levels, + get_ome_zarr_info, +) + + +def _levels(store) -> int: + attributes = dict(zarr.open_group(str(store), mode="r").attrs) + return len(attributes.get("ome", attributes)["multiscales"][0]["datasets"]) + + +def _only_store(root): + return next(root.rglob("*.ome.zarr")) + + +def _geometry() -> Attribute: + attributes = Attribute() + attributes["Spacing"] = np.array([2.0, 1.0, 1.0]) + attributes["Origin"] = np.array([0.0, 0.0, 0.0]) + return attributes + + +def test_a_declared_pyramid_is_written_by_both_paths_and_they_agree(tmp_path): + """``scale_factors`` reaches the store from the dataset, streamed or not, with the same pixels. + + The streamed path cannot take ``scale_factors`` at creation -- no level exists until the last + region lands -- so it derives them at finalize instead. That is different code, and the levels it + produces have to be the same ones, or a chain's output would depend on whether it happened to + stream. + """ + data = np.arange(1 * 16 * 16 * 16, dtype=np.float32).reshape(1, 16, 16, 16) + + Dataset(tmp_path / "whole", "omezarr", scale_factors=[4]).write("G", "case", data, _geometry()) + whole = _only_store(tmp_path / "whole") + + streamed_dataset = Dataset(tmp_path / "streamed", "omezarr", scale_factors=[4]) + stream = streamed_dataset.open_data_stream("G", "case", [1, 16, 16, 16], np.dtype("float32"), _geometry()) + assert stream is not None + with stream: + for start in range(0, 16, 4): + stream.write_slice( + (slice(0, 1), slice(start, start + 4), slice(0, 16), slice(0, 16)), data[:, start : start + 4] + ) + streamed = _only_store(tmp_path / "streamed") + + assert _levels(whole) == 2 + assert _levels(streamed) == 2 + # Level 0 must survive deriving the levels above it: append rewrites the whole multiscales, and an + # earlier version of that truncated the store before dask had pulled a single tile, leaving a + # uniformly zero pyramid with correct metadata and no error. + back, _ = Dataset(tmp_path / "streamed", "omezarr").read_data("G", "case") + assert np.array_equal(np.asarray(back, dtype=np.float32).reshape(data.shape), data) + + coarse_whole, _ = Dataset(tmp_path / "whole", "omezarr@1").read_data("G", "case") + coarse_streamed, _ = Dataset(tmp_path / "streamed", "omezarr@1").read_data("G", "case") + assert np.array_equal(np.asarray(coarse_whole), np.asarray(coarse_streamed)) + # Each level carries its OWN scale; the coarse one is the factor times the fine one. NGFF scale is + # (c, z, y, x) where Spacing is (x, y, z) -- the reversal is the point of asserting it here. + assert get_ome_zarr_info(streamed, 0)["scale"] == [1.0, 1.0, 1.0, 2.0] + assert get_ome_zarr_info(streamed, 1)["scale"] == [1.0, 4.0, 4.0, 8.0] + + +def test_an_interrupted_level_append_leaves_the_original_store_readable(tmp_path, monkeypatch): + """Deriving the coarse levels rewrites level 0, so a failure here is a failure over the only copy. + + The safety is bought with a sibling store and a rename: a rename that does not happen has to leave + the original where its readers expect it, not a gap between two deletes. + """ + data = np.arange(1 * 16 * 16 * 16, dtype=np.float32).reshape(1, 16, 16, 16) + Dataset(tmp_path / "out", "omezarr").write("G", "case", data, _geometry()) + store = _only_store(tmp_path / "out") + real_rename = Path.rename + + def interrupt_the_publishing_rename(self, target): + if self.name.endswith(".appending"): + raise KeyboardInterrupt + return real_rename(self, target) + + monkeypatch.setattr(Path, "rename", interrupt_the_publishing_rename) + with pytest.raises(KeyboardInterrupt): + append_ome_zarr_levels(store, [4]) + monkeypatch.undo() + + back, _ = Dataset(tmp_path / "out", "omezarr").read_data("G", "case") + assert np.array_equal(np.asarray(back, dtype=np.float32).reshape(data.shape), data) + + +def test_a_pyramid_asked_of_a_format_without_levels_is_refused(tmp_path): + """Refused at construction, not ignored: only OME-NGFF has levels, and silently writing one would + leave a consumer's ``@1`` resolving to a level that was never written.""" + with pytest.raises(DatasetManagerError, match="no levels"): + Dataset(tmp_path / "out", "mha", scale_factors=[4]) + + +def test_a_scale_factor_below_two_is_refused(): + """A 'pyramid' whose level does not shrink is a second copy of level 0.""" + from konfai.data.transform import Write + from konfai.utils.errors import TransformError + + with pytest.raises(TransformError, match="scale factor below 2"): + Write(dataset="./Out:omezarr", scale_factors=[1]) diff --git a/tests/unit/test_write_pyramid_and_field_bound.py b/tests/unit/test_write_pyramid_and_field_bound.py deleted file mode 100644 index 7bd3716e..00000000 --- a/tests/unit/test_write_pyramid_and_field_bound.py +++ /dev/null @@ -1,241 +0,0 @@ -# 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 - -"""The write side of the data surface: a declared pyramid, and a field that records its own bound. - -Both are properties a CONSUMER depends on and cannot re-derive cheaply. A pyramid is indexed by -position (``:omezarr@1``), so a producer that writes one level where two were promised does not fail --- it resolves ``@1`` to something else, or to nothing. A field's bound sizes the region a warp must -read, and measuring it after the fact means scanning a volume that exists region by region precisely -because it does not fit. - -So each test here asserts the two write paths agree: assembled in memory, and streamed region by -region. They are different code (``write_ome_zarr`` against ``create_ome_zarr_store`` + -``append_ome_zarr_levels``), and the streamed one is the path a real volume takes. -""" - -from pathlib import Path - -import numpy as np -import pytest -import zarr -from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, Attribute, Dataset -from konfai.utils.errors import DatasetManagerError -from konfai.utils.ome_zarr import ( - DISPLACEMENT_BOUND_ATTRIBUTE, - _read_konfai_attributes, - _zarr_v3_available, - append_ome_zarr_levels, - get_ome_zarr_info, - is_displacement_field, - write_ome_zarr, -) - - -def _levels(store) -> int: - attributes = dict(zarr.open_group(str(store), mode="r").attrs) - return len(attributes.get("ome", attributes)["multiscales"][0]["datasets"]) - - -def _only_store(root): - return next(root.rglob("*.ome.zarr")) - - -def _geometry() -> Attribute: - attributes = Attribute() - attributes["Spacing"] = np.array([2.0, 1.0, 1.0]) - attributes["Origin"] = np.array([0.0, 0.0, 0.0]) - return attributes - - -def test_a_declared_pyramid_is_written_by_both_paths_and_they_agree(tmp_path): - """``scale_factors`` reaches the store from the dataset, streamed or not, with the same pixels. - - The streamed path cannot take ``scale_factors`` at creation -- no level exists until the last - region lands -- so it derives them at finalize instead. That is different code, and the levels it - produces have to be the same ones, or a chain's output would depend on whether it happened to - stream. - """ - data = np.arange(1 * 16 * 16 * 16, dtype=np.float32).reshape(1, 16, 16, 16) - - Dataset(tmp_path / "whole", "omezarr", scale_factors=[4]).write("G", "case", data, _geometry()) - whole = _only_store(tmp_path / "whole") - - streamed_dataset = Dataset(tmp_path / "streamed", "omezarr", scale_factors=[4]) - stream = streamed_dataset.open_data_stream("G", "case", [1, 16, 16, 16], np.dtype("float32"), _geometry()) - assert stream is not None - with stream: - for start in range(0, 16, 4): - stream.write_slice( - (slice(0, 1), slice(start, start + 4), slice(0, 16), slice(0, 16)), data[:, start : start + 4] - ) - streamed = _only_store(tmp_path / "streamed") - - assert _levels(whole) == 2 - assert _levels(streamed) == 2 - # Level 0 must survive deriving the levels above it: append rewrites the whole multiscales, and an - # earlier version of that truncated the store before dask had pulled a single tile, leaving a - # uniformly zero pyramid with correct metadata and no error. - back, _ = Dataset(tmp_path / "streamed", "omezarr").read_data("G", "case") - assert np.array_equal(np.asarray(back, dtype=np.float32).reshape(data.shape), data) - - coarse_whole, _ = Dataset(tmp_path / "whole", "omezarr@1").read_data("G", "case") - coarse_streamed, _ = Dataset(tmp_path / "streamed", "omezarr@1").read_data("G", "case") - assert np.array_equal(np.asarray(coarse_whole), np.asarray(coarse_streamed)) - # Each level carries its OWN scale; the coarse one is the factor times the fine one. NGFF scale is - # (c, z, y, x) where Spacing is (x, y, z) -- the reversal is the point of asserting it here. - assert get_ome_zarr_info(streamed, 0)["scale"] == [1.0, 1.0, 1.0, 2.0] - assert get_ome_zarr_info(streamed, 1)["scale"] == [1.0, 4.0, 4.0, 8.0] - - -def test_an_interrupted_level_append_leaves_the_original_store_readable(tmp_path, monkeypatch): - """Deriving the coarse levels rewrites level 0, so a failure here is a failure over the only copy. - - The safety is bought with a sibling store and a rename: a rename that does not happen has to leave - the original where its readers expect it, not a gap between two deletes. - """ - data = np.arange(1 * 16 * 16 * 16, dtype=np.float32).reshape(1, 16, 16, 16) - Dataset(tmp_path / "out", "omezarr").write("G", "case", data, _geometry()) - store = _only_store(tmp_path / "out") - real_rename = Path.rename - - def interrupt_the_publishing_rename(self, target): - if self.name.endswith(".appending"): - raise KeyboardInterrupt - return real_rename(self, target) - - monkeypatch.setattr(Path, "rename", interrupt_the_publishing_rename) - with pytest.raises(KeyboardInterrupt): - append_ome_zarr_levels(store, [4]) - monkeypatch.undo() - - back, _ = Dataset(tmp_path / "out", "omezarr").read_data("G", "case") - assert np.array_equal(np.asarray(back, dtype=np.float32).reshape(data.shape), data) - - -def test_a_pyramid_asked_of_a_format_without_levels_is_refused(tmp_path): - """Refused at construction, not ignored: only OME-NGFF has levels, and silently writing one would - leave a consumer's ``@1`` resolving to a level that was never written.""" - with pytest.raises(DatasetManagerError, match="no levels"): - Dataset(tmp_path / "out", "mha", scale_factors=[4]) - - -def test_a_scale_factor_below_two_is_refused(): - """A 'pyramid' whose level does not shrink is a second copy of level 0.""" - from konfai.data.transform import Write - from konfai.utils.errors import TransformError - - with pytest.raises(TransformError, match="scale factor below 2"): - Write(dataset="./Out:omezarr", scale_factors=[1]) - - -# An RFC-5 field is a zarr v3 store (NGFF >= 0.5), which zarr 2.x -- the newest release for Python -# 3.10 -- cannot write. The pyramid tests above need no such store and run everywhere. -_needs_rfc5 = pytest.mark.skipif( - not _zarr_v3_available(), - reason="a displacement field's bound is recorded in a zarr v3 store (zarr>=3, Python>=3.11)", -) - - -@_needs_rfc5 -def test_a_field_records_its_own_bound_on_both_write_paths(tmp_path): - """The largest displacement per component, in world units, written by whoever holds the samples. - - Per component and not one scalar: a consumer divides each bound by the spacing of its OWN axis, - and these grids are anisotropic, so one collapsed maximum over-reads two axes out of three. - """ - rng = np.random.default_rng(0) - field = rng.normal(0.0, 30.0, size=(3, 8, 8, 8)).astype(np.float32) - field[0, 4, 4, 4], field[1, 2, 2, 2] = 917.5, -640.25 - expected = [float(np.abs(field[component]).max()) for component in range(3)] - - whole = tmp_path / "whole" / "case" / "DVF.ome.zarr" - whole.parent.mkdir(parents=True) - write_ome_zarr(whole, field, spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0), displacement_field=True) - assert _read_konfai_attributes(whole)[DISPLACEMENT_BOUND_ATTRIBUTE] == expected - - shape, attributes = Dataset(tmp_path / "whole", "omezarr").get_infos("DVF", "case") - # The read side must carry the RFC-5 type on the HEADERS path too, or a field read region by - # region and written region by region comes out an ordinary 3-channel image. - assert DISPLACEMENT_FIELD_ATTRIBUTE in attributes - - stream = Dataset(tmp_path / "streamed", "omezarr").open_data_stream( - "DVF", "case", list(shape), np.dtype("float32"), attributes - ) - assert stream is not None - with stream: - for start in range(0, 8, 2): - stream.write_slice( - (slice(0, 3), slice(start, start + 2), slice(0, 8), slice(0, 8)), field[:, start : start + 2] - ) - streamed = _only_store(tmp_path / "streamed") - - # Accumulated across regions: the streamed path never sees the field whole, which is the whole - # reason the bound has to be recorded rather than measured later. - assert _read_konfai_attributes(streamed)[DISPLACEMENT_BOUND_ATTRIBUTE] == expected - assert is_displacement_field(streamed) - assert DISPLACEMENT_BOUND_ATTRIBUTE in Dataset(tmp_path / "streamed", "omezarr").get_infos("DVF", "case")[1] - - -def tmp_field_store(field: np.ndarray) -> Path: - import tempfile - - root = Path(tempfile.mkdtemp()) - store = root / "fields" / "case" / "DVF.ome.zarr" - store.parent.mkdir(parents=True) - write_ome_zarr(store, field, spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0), displacement_field=True) - return root / "fields" - - -@_needs_rfc5 -def test_the_recorded_bound_reaches_each_axis_by_its_own_spacing_under_anisotropy() -> None: - """What the bound is FOR, read by the stage that consumes it. - - Component ``i`` of a displacement field is world axis (x, y, z)[i], while array axes are - (z, y, x). So the reach of a region on each array axis reads the components reversed, each - against its own spacing. Getting the pairing wrong is a warp that raises nothing and reads the - wrong neighbourhood -- and under this anisotropy the three numbers are far enough apart (3, 22 - and 31 voxels) that any permutation of them is visible. - """ - from konfai.data.transform import LocalityKind, Resample - - field = np.zeros((3, 8, 8, 8), dtype=np.float32) - field[0, 4, 4, 4], field[1, 2, 2, 2], field[2, 1, 1, 1] = 917.5, -640.25, 96.0 - store = tmp_field_store(field) - - warp = Resample(field=f"{store}:omezarr", field_group="DVF") - attribute = Attribute() - attribute["Spacing"] = np.array([30.08, 30.08, 40.0]) # stored (x, y, z) - attribute["Origin"] = np.zeros(3) - attribute["Direction"] = np.eye(3).reshape(-1) - - assert warp.patch_locality(attribute).kind is LocalityKind.REGRID - - shape = [64, 128, 128] - warp.transform_shape("CT", "CASE_000", shape, attribute) - target = tuple(slice(30, 32) for _ in shape) - window = warp.stream_region_source("CASE_000", target, shape, attribute) - - # z takes the z component (96 um over a 40 um voxel), x the x component (917.5 over 30.08). - per_axis = [(96.0, 40.0), (640.25, 30.08), (917.5, 30.08)] # array order (z, y, x) - expected = [ - ( - max(0, int(np.floor(30 - 0.5 - reach / spacing)) - 1), - min(extent, int(np.ceil(32 - 0.5 + reach / spacing)) + 2), - ) - for (reach, spacing), extent in zip(per_axis, shape, strict=True) - ] - assert [(part.start, part.stop) for part in window] == expected From 4b0f669cff0cc9872378a22d9b874d51c1224caf Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 04:27:16 +0200 Subject: [PATCH 30/39] feat(data): the itktransform backend reads by regions The parameters are HDF5, so a span of leading-axis rows is one contiguous span of the dataset: a region of a displacement entry decodes the rows it maps to and nothing else, and bounded_region_reads says so to the plan. A foreign or non-displacement file keeps the whole-read path. --- konfai/utils/dataset.py | 32 ++++++++++++++++++++++-- tests/unit/test_itk_transform_backend.py | 13 ++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index ad4fffc3..41e662cb 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -1950,9 +1950,37 @@ def file_to_data(self, group: str, name: str) -> tuple[np.ndarray, Attribute]: attributes, ) + def bounded_region_reads(self, name: str) -> bool: + try: + import h5py # noqa: F401 + except ImportError: + return False + shape, _attributes = self.get_infos("", name) + return len(shape) == 4 and shape[0] == 3 + def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: - data, attributes = self.file_to_data(group, name) - return data[slices], attributes + """A region of a displacement entry, decoded from the parameters it maps to alone. + + The buffer is ``[z][y][x]`` with the component fastest, so a span of leading-axis rows + is one contiguous span of the parameters: read, reshaped, and sliced down to the exact + region — the peak is the row span, never the field. + """ + try: + import h5py + except ImportError: + data, attributes = self.file_to_data(group, name) + return data[slices], attributes + shape, attributes = self.get_infos(group, name) + if len(shape) != 4 or shape[0] != 3 or DISPLACEMENT_FIELD_ATTRIBUTE not in attributes: + data, attributes = self.file_to_data(group, name) + return data[slices], attributes + spatial = shape[1:] + leading = slices[1].indices(spatial[0]) + row = 3 * int(np.prod(spatial[1:], dtype=np.int64)) + with h5py.File(self._path(name), "r") as file: + span = file["TransformGroup/0/TransformParameters"][leading[0] * row : leading[1] * row] + block = np.moveaxis(span.reshape(leading[1] - leading[0], *spatial[1:], 3), -1, 0) + return np.asarray(block[(slices[0], slice(None), *slices[2:])], dtype=np.float32), attributes def file_to_data_statistics( self, diff --git a/tests/unit/test_itk_transform_backend.py b/tests/unit/test_itk_transform_backend.py index 096cd84d..216f5781 100644 --- a/tests/unit/test_itk_transform_backend.py +++ b/tests/unit/test_itk_transform_backend.py @@ -122,3 +122,16 @@ def test_a_foreign_affine_file_reads_back_too(tmp_path: Path) -> None: back = Dataset(tmp_path / "out", "itktransform").read_transform("Reg", "P000") point = (1.0, 2.0, 3.0) assert back.TransformPoint(point) == pytest.approx(affine.TransformPoint(point)) + + +def test_a_region_read_decodes_only_its_rows_and_matches_the_whole(tmp_path: Path) -> None: + """The parameters are HDF5, so a slab reads the span it maps to — same values as the whole.""" + field = _field(3) + dataset = Dataset(tmp_path / "out", "itktransform") + dataset.write("Transform", "P000", field, _attributes()) + + assert dataset.bounded_region_reads("Transform", "P000") + whole, _ = dataset.read_data("Transform", "P000") + region, _ = dataset.read_data_slice("Transform", "P000", (slice(0, 3), slice(1, 3), slice(2, 5), slice(0, 4))) + + np.testing.assert_array_equal(np.asarray(region), np.asarray(whole)[:, 1:3, 2:5, 0:4]) From 3210e809bd68963d70e367e5961368f4943e8698 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 04:28:16 +0200 Subject: [PATCH 31/39] docs(data): the backend table says what streams, and the transform backend joins it --- .../reference/components/storage-backends.md | 34 +++++++++++++++---- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/docs/source/reference/components/storage-backends.md b/docs/source/reference/components/storage-backends.md index 521ed5b8..0b8eff55 100644 --- a/docs/source/reference/components/storage-backends.md +++ b/docs/source/reference/components/storage-backends.md @@ -9,21 +9,41 @@ serve; the DICOM and OME-Zarr reader APIs are detailed below. ## Backends -| Backend | Format token(s) | Kind | Optional extra | -| --- | --- | --- | --- | -| `Dataset.SitkFile` | `mha, mhd, nii, nii.gz, nrrd, nrrd.gz, gipl(.gz), hdr, img, dcm, tif(f), png, jpg, jpeg, bmp, itk.txt, fcsv, xml, vtk, npy` | Directory of per-case image files (default) | `konfai[itk]` (`SimpleITK`) | -| `Dataset.H5File` | `h5` | Single monolithic HDF5 file | `konfai[hdf5]` (`h5py`) | -| `Dataset.OmeZarrFile` | `omezarr, ome-zarr, ome_zarr, zarr` (+ `@level`) | OME-Zarr pyramid directory | `konfai[omezarr]` (`zarr` + `ngff-zarr`) | -| `Dataset.DicomFile` (DICOM series; scalar-array writes) | `dicom` | DICOM series directory | `konfai[dicom]` (`pydicom`) | +The two streaming columns are what the planner prices: **region reads** says +whether a region decodes only itself (a backend that answers "no" decodes the +whole volume behind every region, which only ever costs speed, never +correctness), and **streamed writes** says whether `open_data_stream` can build +the entry region by region (otherwise the volume is assembled and written +whole). + +| Backend | Format token(s) | Kind | Region reads | Streamed writes | Optional extra | +| --- | --- | --- | --- | --- | --- | +| `Dataset.SitkFile` | `mha, mhd, nii, nii.gz, nrrd, nrrd.gz, gipl(.gz), hdr, img, dcm, tif(f), png, jpg, jpeg, bmp, itk.txt, fcsv, xml, vtk, npy` | Directory of per-case image files (default) | **uncompressed MetaImage and NIfTI only** — compressed streams are not seekable, and NRRD never streams in ITK | **`.mha` only** (memmap over the raw pixel block; needs image geometry) | `konfai[itk]` (`SimpleITK`) | +| `Dataset.H5File` | `h5` | Single monolithic HDF5 file | yes (chunked) | yes | `konfai[hdf5]` (`h5py`) | +| `Dataset.OmeZarrFile` | `omezarr, ome-zarr, ome_zarr, zarr` (+ `@level`) | OME-Zarr pyramid directory | yes (chunked) | yes, `scale_factors` pyramids included | `konfai[omezarr]` (`zarr` + `ngff-zarr`) | +| `Dataset.DicomFile` (DICOM series; scalar-array writes) | `dicom` | DICOM series directory | per slice | no (whole series) | `konfai[dicom]` (`pydicom`) | +| `Dataset.ItkTransformFile` | `itktransform` | ITK transform files (`.h5`, `.tfm`), one per case/group | yes for a displacement entry (the parameters are HDF5; a row span is one contiguous read) | yes — the parameters fill region by region, and the file is what `sitk.WriteTransform` would have written | `konfai[itk]`; `h5py` for the region paths (whole-file sitk fallback without it) | ```{tip} -`pip install "konfai[imaging]"` installs **all four** backends at once +`pip install "konfai[imaging]"` installs every backend at once (`SimpleITK, h5py, pydicom, zarr, ngff-zarr`). ``` The extras column above is the summary; {doc}`../../getting-started/installation` is the canonical home for the optional-extras table. +## ITK transform files as a dataset + +`:itktransform` stores one ITK transform per entry (`/.h5`). The +write side is the point: a displacement field streams into the exact file +`sitk.WriteTransform` would produce — pinned identical through ITK's own reader +— without ever holding the field whole in float64. The read side hands back what +`Dataset.read_transform` decodes: a displacement entry as its field (region +reads included), any other stored transform (affine, composite, `.tfm`) as its +parameters. This is how a run's `Transform.h5` deliverable is a plain `Write:` +like any other, and how a staged transform file resolves through the same +`Dataset` surface as an image. + ## The `SitkFile` default backend also handles sidecars Beyond images, the SITK backend reads/writes several sidecar payloads by From 20489a9e3d85bc346056c7b4289ea635d34708dd Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 04:37:25 +0200 Subject: [PATCH 32/39] feat(data): stream uncompressed NIfTI writes, like the reads The region-writable SimpleITK set becomes the region-readable one: uncompressed .nii is a fixed 348-byte header plus a flat raw block, exactly the property the .mha stream memmaps -- and NIfTI's vector dimension is its slowest, so channel-first slabs land without a transpose. The one convention the stream owns is the RAS sform (the pipeline speaks LPS; the affine's first two rows negate on the way out), pinned against sitk's own writer on an oblique grid. nii joins the stream test matrix, so concurrency, abort and invisibility hold for it as for the others. --- docs/source/config_guide/transform.md | 2 +- .../reference/components/storage-backends.md | 2 +- konfai/utils/dataset.py | 97 +++++++++++++++++-- tests/unit/test_data_stream.py | 47 ++++++++- 4 files changed, 137 insertions(+), 11 deletions(-) diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index 6fc9c690..ee41c05c 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -116,7 +116,7 @@ The usual reasons a chain refuses to stream, and what to do: | --- | --- | | a stage declares `WHOLE_VOLUME` | Some transforms genuinely need the volume (`Squeeze`, `Norm` change the tensor's rank). Nothing to fix — check it fits the budget. | | a statistic after a value-changing stage | Insert a `Save:` before the statistic; the cache becomes the source the statistic reads. See below. | -| the destination cannot serve region writes | Write to `:h5` or `:omezarr` (or `:mha` for an image with geometry). | +| the destination cannot serve region writes | Write to `:h5` or `:omezarr` (or `:mha`/`:nii` for an image with geometry). | | a halo too wide for the grid | The transform's neighbourhood is over half the slab extent; it is cheaper to load the volume. | ### `Save:` unlocks chains that would otherwise refuse diff --git a/docs/source/reference/components/storage-backends.md b/docs/source/reference/components/storage-backends.md index 0b8eff55..e07e6a22 100644 --- a/docs/source/reference/components/storage-backends.md +++ b/docs/source/reference/components/storage-backends.md @@ -18,7 +18,7 @@ whole). | Backend | Format token(s) | Kind | Region reads | Streamed writes | Optional extra | | --- | --- | --- | --- | --- | --- | -| `Dataset.SitkFile` | `mha, mhd, nii, nii.gz, nrrd, nrrd.gz, gipl(.gz), hdr, img, dcm, tif(f), png, jpg, jpeg, bmp, itk.txt, fcsv, xml, vtk, npy` | Directory of per-case image files (default) | **uncompressed MetaImage and NIfTI only** — compressed streams are not seekable, and NRRD never streams in ITK | **`.mha` only** (memmap over the raw pixel block; needs image geometry) | `konfai[itk]` (`SimpleITK`) | +| `Dataset.SitkFile` | `mha, mhd, nii, nii.gz, nrrd, nrrd.gz, gipl(.gz), hdr, img, dcm, tif(f), png, jpg, jpeg, bmp, itk.txt, fcsv, xml, vtk, npy` | Directory of per-case image files (default) | **uncompressed MetaImage and NIfTI only** — compressed streams are not seekable, and NRRD never streams in ITK | **uncompressed `.mha` and `.nii`** (memmap over the raw pixel block; needs image geometry) — the region-writable set is the region-readable one, deliberately | `konfai[itk]` (`SimpleITK`) | | `Dataset.H5File` | `h5` | Single monolithic HDF5 file | yes (chunked) | yes | `konfai[hdf5]` (`h5py`) | | `Dataset.OmeZarrFile` | `omezarr, ome-zarr, ome_zarr, zarr` (+ `@level`) | OME-Zarr pyramid directory | yes (chunked) | yes, `scale_factors` pyramids included | `konfai[omezarr]` (`zarr` + `ngff-zarr`) | | `Dataset.DicomFile` (DICOM series; scalar-array writes) | `dicom` | DICOM series directory | per slice | no (whole series) | `konfai[dicom]` (`pydicom`) | diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 41e662cb..5e321bac 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -790,6 +790,80 @@ def _close(self, success: bool) -> None: Path(self._temporary_path).unlink(missing_ok=True) +# NIfTI-1 datatype code for each NumPy dtype a streamed .nii can hold. +_NIFTI_DATATYPES = { + "uint8": 2, + "int16": 4, + "int32": 8, + "float32": 16, + "float64": 64, + "int8": 256, + "uint16": 512, + "uint32": 768, + "int64": 1024, + "uint64": 1280, +} + + +class _NiftiDataStream(DataStream): + """Uncompressed NIfTI-1 written region by region: a hand-written 348-byte header, then a memmap + over the raw block. NIfTI's data order is x fastest with the vector dimension SLOWEST, which is + exactly the channel-first ``[C, Z, Y, X]`` layout in C order — the map is the block itself. + The sform carries the geometry, and NIfTI speaks RAS where the pipeline speaks LPS: the + affine's first two rows are negated on the way out, the one convention this class owns.""" + + def __init__(self, path: str, shape: list[int], dtype: np.dtype, attributes: Attribute) -> None: + import struct + + self.path = path + self._temporary_path = f"{path}.{self.temporary_suffix()}" + channels, spatial = int(shape[0]), [int(extent) for extent in shape[1:]] + # The header is written little-endian, so the block must be too. + self._dtype = np.dtype(dtype).newbyteorder("<") + size_xyz = spatial[::-1] + spacing = attributes.get_np_array("Spacing").astype(np.float64) + origin = attributes.get_np_array("Origin").astype(np.float64) + direction = attributes.get_np_array("Direction").astype(np.float64).reshape(3, 3) + affine = np.concatenate([direction * spacing[np.newaxis, :], origin[:, np.newaxis]], axis=1) + affine[:2] *= -1.0 # LPS -> RAS + header = bytearray(348) + struct.pack_into(" 1: + struct.pack_into(" None: + self._memmap[slices] = data + + def _close(self, success: bool) -> None: + self._memmap.flush() + del self._memmap + if success: + os.replace(self._temporary_path, self.path) + else: + os.remove(self._temporary_path) + + # MetaImage ElementType for each NumPy dtype a streamed .mha can hold. _MHA_ELEMENT_TYPES = { "int8": "MET_CHAR", @@ -1522,21 +1596,28 @@ def open_data_stream( attributes: Attribute, region_shape: list[int] | None = None, ) -> DataStream | None: - # Only an uncompressed local-data MetaImage is region-writable (ASCII header + flat raw - # block); every other SimpleITK format writes the whole image in one WriteImage call. - if self.file_format != "mha" or not is_an_image(attributes) or len(shape) < 3: + # The region-writable SimpleITK formats are the region-READABLE ones, deliberately: + # uncompressed MetaImage and NIfTI are a fixed header plus a flat raw block, so the block + # is reserved and memmapped. Every other format writes whole in one WriteImage call -- + # and streaming into a form the reader must then decode whole would only move the cost. + if self.file_format not in ("mha", "nii") or not is_an_image(attributes) or len(shape) < 3: return None element_dtype = np.dtype(dtype) if element_dtype == np.float16: - # MetaImage has no half-float type; widen float16 to float32 (exact), as data_to_image - # does, so streamed and whole-volume writes hold identical bytes. + # Neither format has a half-float type; widen float16 to float32 (exact), as + # data_to_image does, so streamed and whole-volume writes hold identical bytes. element_dtype = np.dtype(np.float32) - if element_dtype.name not in _MHA_ELEMENT_TYPES: - return None dimension = len(shape) - 1 geometry = (("Origin", dimension), ("Spacing", dimension), ("Direction", dimension * dimension)) if any(len(attributes.get_np_array(key)) != n for key, n in geometry): return None + if self.file_format == "nii": + if dimension != 3 or element_dtype.name not in _NIFTI_DATATYPES: + return None + os.makedirs(self.filename, exist_ok=True) + return _NiftiDataStream(f"{self.filename}{name}.{self.file_format}", shape, element_dtype, attributes) + if element_dtype.name not in _MHA_ELEMENT_TYPES: + return None os.makedirs(self.filename, exist_ok=True) return _MhaDataStream(f"{self.filename}{name}.{self.file_format}", shape, element_dtype, attributes) @@ -2275,7 +2356,7 @@ def can_stream_data(self, attributes: Attribute) -> bool: return True if self.file_format == "itktransform": return is_an_image(attributes) - return self.file_format == "mha" and is_an_image(attributes) + return self.file_format in ("mha", "nii") and is_an_image(attributes) def open_data_stream( self, diff --git a/tests/unit/test_data_stream.py b/tests/unit/test_data_stream.py index 7a5f4002..c19cc378 100644 --- a/tests/unit/test_data_stream.py +++ b/tests/unit/test_data_stream.py @@ -56,7 +56,7 @@ def _skip_unavailable(file_format: str) -> None: pytest.importorskip("h5py") -FORMATS = ["mha", "h5", "omezarr"] +FORMATS = ["mha", "nii", "h5", "omezarr"] @pytest.mark.parametrize("file_format", FORMATS) @@ -295,3 +295,48 @@ def test_can_stream_data_matches_open_support(tmp_path: Path) -> None: assert not Dataset(tmp_path / "b", "nii.gz").can_stream_data(geometry) assert Dataset(tmp_path / "c", "h5").can_stream_data(Attribute()) assert Dataset(tmp_path / "d", "omezarr").can_stream_data(Attribute()) + + +def test_nii_stream_is_the_file_sitk_would_have_written(tmp_path: Path) -> None: + """The one convention the NIfTI stream owns is the RAS sform; sitk's own writer is the oracle. + + Compared through sitk's reader on an OBLIQUE grid: a dropped or half-applied LPS-to-RAS flip + reads back as a different Origin/Direction, not as an error. + """ + import SimpleITK as sitk + + volume = _volume(channels=1) + attributes = _image_attributes() + angle = np.deg2rad(30.0) + cos, sin = float(np.cos(angle)), float(np.sin(angle)) + attributes["Direction"] = np.asarray([[cos, -sin, 0.0], [sin, cos, 0.0], [0.0, 0.0, 1.0]]).reshape(-1) + + dataset = Dataset(tmp_path / "streamed", "nii") + _write_by_slabs(dataset, volume, attributes) + + reference = sitk.GetImageFromArray(volume[0]) + reference.SetOrigin(attributes.get_np_array("Origin").tolist()) + reference.SetSpacing(attributes.get_np_array("Spacing").tolist()) + reference.SetDirection(attributes.get_np_array("Direction").tolist()) + sitk.WriteImage(reference, str(tmp_path / "reference.nii")) + + got = sitk.ReadImage(str(tmp_path / "streamed" / "CASE_001" / "CT.nii")) + want = sitk.ReadImage(str(tmp_path / "reference.nii")) + np.testing.assert_array_equal(sitk.GetArrayFromImage(got), sitk.GetArrayFromImage(want)) + np.testing.assert_allclose(got.GetOrigin(), want.GetOrigin(), atol=1e-5) + np.testing.assert_allclose(got.GetSpacing(), want.GetSpacing(), atol=1e-6) + np.testing.assert_allclose(got.GetDirection(), want.GetDirection(), atol=1e-6) + + +def test_nii_stream_multi_channel_reads_back_as_vector_image(tmp_path: Path) -> None: + """The vector dimension is NIfTI's slowest, so channel-first slabs land without a transpose.""" + import SimpleITK as sitk + + volume = _volume(channels=3) + dataset = Dataset(tmp_path / "streamed", "nii") + _write_by_slabs(dataset, volume, _image_attributes()) + + image = sitk.ReadImage(str(tmp_path / "streamed" / "CASE_001" / "CT.nii")) + assert image.GetNumberOfComponentsPerPixel() == 3 + back, _ = dataset.read_data("CT", "CASE_001") + np.testing.assert_array_equal(np.asarray(back), volume) From 2f6171880352cff94347ff2486cd4c70b6fbfeca Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 12:24:31 +0200 Subject: [PATCH 33/39] docs: the Python page; the changelog is written at the tag Drop the hand-written Unreleased section and the entries the branch had appended into the released v1.8.0 section: v1.8.0 shipped without this branch, and the file's own preamble says each section is drafted by commitizen at tag time, then edited. The next tag documents this line. --- CHANGELOG.md | 138 ------------------ docs/source/config_guide/transform.md | 7 + .../source/reference/components/transforms.md | 5 +- docs/source/usage/index.rst | 1 + docs/source/usage/python-workflows.md | 60 ++++++++ 5 files changed, 71 insertions(+), 140 deletions(-) create mode 100644 docs/source/usage/python-workflows.md diff --git a/CHANGELOG.md b/CHANGELOG.md index ad575c0e..fa86fbde 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,125 +16,6 @@ draft, then say what a user of the package gets that they did not have -- and re against the commits that landed *after* you drafted it. Running the command over a section already written replaces it. -## Unreleased - -### ✨ Features - -- **transform**: one `Resample`. `ResampleToResolution`, `ResampleToShape`, `ResampleToReference`, - `ResampleTransform` and `Warp` were five stages answering two questions between them, each with a - sampler of its own. They are now five spellings of one stage that asks the questions separately: - **which grid to write on** (nothing, `spacing`, `shape`, or `reference` — a stored image's grid - adopted whole) and **what map to write it through** (`field`, `transforms`, or neither). Every - combination is legal, and asked for together they compose into **one interpolation** instead of - two. **The old names are gone**, so a config naming one must be migrated: the class becomes - `Resample`, and `ResampleToReference`'s `entry` / `group` / `dataset` become `reference` / - `reference_group` / `reference_dataset`. Every other argument carries over unchanged. -- **transform**: `align` says where a `spacing` or `shape` grid sits — `extent` (the default) keeps - the field of view, `origin` keeps voxel zero's centre. This was decided silently before, and - differently by the data and by the header. -- **transform**: a resample streams whatever it is asked for, and on the volume's device. Applying a - stored registration used to hold a whole volume — not because a warp needs one, but because - nothing on a `sitk.Transform` said how far it reached. A rigid or affine map now bounds exactly; a - BSpline and a displacement field bound by the largest of their values, which holds at every point - rather than at the sampled ones. Everything is evaluated in torch, so nothing marshals a - GPU-resident case out to numpy and back. -- **transform**: a resample no longer requires the grids to share a direction. A rotated reference, - or a field stored on turned axes, used to be refused with an instruction to run `Canonical` first - — a second interpolation of the same voxels. Both now work directly. -- **transform**: a field is read on its own grid. `Warp` required the field and the case to share - one; a field solved at 120 µm now moves a volume stored at 30 µm without being upsampled first. -- **transform**: `ResampleToShape` needs no geometry at all. A count is a count; only a change of - density needs the density it starts from. - -- **transform**: the plan chooses the route from predicted cost against the memory budget. - Streaming is a memory strategy — splitting re-reads, loading whole reads once — so a case that - fits the budget is now `LOAD`ed when streaming would re-read the source (a halo re-reads its - overlap, a regrid pulls each slab's window through its map, a compressed store decodes the whole - volume per slab — measured up to 9.7x on an oblique map). The plan prints the predicted factor; - `on_fallback` has nothing to say about a choice. The predictor's streamed-vs-assembled route now - prices the config's budget instead of the machine's free memory at that moment, so the same case - takes the same route on a loaded machine and an idle one. -- **transform**: a console that says what changes a decision. A healthy 6-output run is 7 lines — - the plan's chain lines now carry the stages and the terminal `Write` destination, and one final - line states what was written, how, in how long, and where `outputs.json` is. The run itself only - speaks when it deviates from the printed plan. A designed refusal (`on_fallback: error`, a budget - overrun) prints its message and remedy and exits 1 — 34 lines of framework traceback down to 6; - `KONFAI_DEBUG=1` re-attaches the traceback. Logs stop amplifying progress frames (~4x smaller), - and every byte figure prints at the unit that carries digits instead of `0.00 GiB`. -- **transform**: every configuration-dependent fallback says what to change — a masked `Clip` or - `Standardize`, a percentile bound, a spatial `Sum`/`Argmax`/`Softmax`, an oblique `Canonical`, a - free-angle `Rotate` draw, a vector-field `Flip` — and `Statistics` streams: its four numbers are - the disk scan's own, seeded instead of recomputed per region. - -### 🐛 Fixes - -- **transform**: a resampled label map no longer comes out shifted against the image beside it. - `F.interpolate`'s nearest reads `floor(o * scale)` where its linear reads `scale * (o + 0.5) - 0.5`, - so a mask resampled by the same stage as its CT lagged it by `(scale - 1) / 2` source voxels — - 2.5 voxels, 1.25 mm of anatomy, resampling 0.5 mm to 3 mm. Both volumes were entirely plausible on - their own. Nearest is now ITK's round-half-up on the same physical index the linear sampler reads. -- **transform**: the header a resample records now describes the grid it actually sampled. - `ResampleToResolution` wrote the spacing that was *asked for* while sampling at `n_in/n_out` times - the source's (up to a millimetre of drift across a volume) and left the `Origin` alone while - sampling half a spacing-change away from it. Nothing downstream could see either: the voxels are - all real, and the header was the only witness. -- **transform**: a voxel count no longer loses a slice to floating point. 90 voxels of 0.7 mm re-cut - at 1.5 mm is 42, and the count went through float32 to get there — landing on 41 or 42 depending - on the numbers. -- **transform**: a warp on an oblique case reads the neighbourhood it needs. The halo was derived - per array axis from a world displacement, which assumes the direction cosines are the identity; on - a turned case the window was short on the axes the displacement actually reached, and a short - window returns the border value rather than raising. -- **transform**: a resample refuses what it used to do quietly. A refusal the whole-volume path can - serve — an undeclared field bound, a case with no geometry — declares `WHOLE_VOLUME` with the - sentence saying what to change, and the run proceeds assembled. A map neither route can apply — - an unsupported spline order, a missing entry, `invert: true` on a spline or a field — refuses as - the plan is built, before a byte is written; it used to print a fallback the run then contradicted - by dying per case. -- **transform**: `ResampleTransform`'s `inverse` defaults to `false`. It always raised - `NotImplementedError`, so a prediction finalize through this stage failed at the end of the run - rather than at its configuration. - -- **data**: a stage is judged on the state the stages before it left, on every landing fold. A - `Resample` behind a `Canonical` recorded the pre-reorientation grid and resampled the wrong axis - — silently, every voxel real, on the exact chain the published TotalSegmentator bundle ships — - and a second `Resample` saw the original spacing and handed its input through as a no-op. Both - now stream, bit-identical to the whole-volume pass. -- **data**: the end plane of a BSpline's valid region is warped as ITK warps it. A grid - commensurate with the coefficient mesh — what a fitted transform domain produces — hit that - plane whole planes at a time, every voxel silently unmoved (2.66 mm of displacement dropped, - against 1e-14 agreement everywhere else). -- **data**: coverage is judged through the declared map before a case is refused as disjoint. An MR - and a CT 1000 mm apart in stage coordinates with a stored rigid bridging them — the situation the - apply step exists to serve — were refused as writing nothing but fill. -- **data**: a half-precision volume on CUDA blends through float32 coordinates. The fused blend - built its sampling grid in the payload's dtype, quantizing a coordinate at ~2^-11 of the window — - 0.06 voxel on a 512 axis, tens of units at a sharp edge (measured 60.0 on a 1000-range fixture, - 0.022 after). -- **data**: interleaved patch reads of two `Expand` copies each keep their own grids; re-reading a - copy after another was planned handed it the other copy's sampling. -- **transform**: `Clip('min'/'max')` clips to the case's seeded statistic, not the region's own — - what `save_clip_min`/`save_clip_max` recorded used to depend on which patch happened to run. - -### ⚡ Performance - -- **data**: a map that factorises is read one axis at a time on global coordinates — most - resamples, bit-identical streamed or whole (CT-sized case: CPU 2269 -> 160 ms, GPU 46.5 -> 2.1 ms, - peak 1.09 -> 0.35 GiB) — and the axis that shrinks most is blended first (9.1x on a thick-slice CT - brought to isotropic). A map that does not factorise (a warp, a rotation, a stored field) goes - through one fused `grid_sample` kernel (4x on a warp); on that path a streamed region and the - whole volume agree to ~1e-5 of the data's range rather than bit for bit, which the plan notes - when a budget shrinks the slabs. - -### 🔧 Internals - -- **data**: `LocalityKind.RESCALE` is gone. It was the dispatcher's own resample map — a size ratio, - which says nothing once a target grid has an origin — and with one resample stage there is one - regime, `REGRID`, that the stage owns both halves of. -- **data**: the sampler owns its rules (`nearest_index`, `window_index`, `sampling_dtype` live in - `sampling.py`), `utils/ITK.py` keeps only its live decoders (~340 orphaned pre-unification lines - deleted), and `KONFAI_STREAM_LINEAR_RESAMPLE` — documented, read by nothing — is out of the docs. - ## v1.8.0 (2026-08-04) ### ✨ Features @@ -150,25 +31,6 @@ written replaces it. - **data**: Vote, the reduction operator that folds segmentations without inventing a label - **data**: declare an OME-Zarr pyramid from a Write, and let a field carry its own bound - **impact-reg**: seed the rigid from the centre of mass, not only the frame -- **impact-reg**: `--tmp-dir` on register/eval/uncertainty, the option the other app CLIs already carry — - a caller whose system temp directory is a tmpfs can now stage volume-sized intermediates on real disk - instead of overriding `TMPDIR` from outside; the same change also writes the moved image and the - displacement field once per run instead of twice -- **impact-reg**: `register` and `eval` accept a whole dataset per input, not only one volume — a directory is - expanded into one case per volume it holds, exactly as `konfai-apps infer` already does, and every case - gets its own field, moved image and transform. Previously the cases were counted from the command-line - arguments while konfai-apps counted them from the expanded units, so a directory input produced N - results and only the first was collected, silently -- **impact-reg**: `register --fields-only` writes the displacement fields and stops there. The moved - image and `Transform.h5` are both derived from the field, at the cost of a full-size resample and a - full-size rewrite, so a caller that composes the field with its own and derives its own moved — the - ExaSPIM tiled refinement does exactly that — no longer pays for two outputs it deletes -- **impact-reg**: a registration preset now owes exactly one output, its displacement field, in whatever - format it declares — `register` derives the moved image from it instead of expecting a second output. - A preset can drop `MovedImage` entirely, which for a tiled one also drops blending a full-size moved - across every patch seam for a caller that has the field. Reading the moving image handles an OME-Zarr - store as well as an ITK file, which fixes the ensemble path too: averaging several presets over - OME-Zarr inputs failed there, and nowhere else, on `sitk.ReadImage` - **studio**: bundle icons through the app interface, and a way to stop Studio (#75) - **examples**: a Transform example -- a template folded out of a cohort, and drawn copies of a case diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index ee41c05c..4d0c7513 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -304,6 +304,13 @@ placement and ignores this. ### `Resample: {reference: …}`: making `strict` true rather than waived +A reference can also **follow the case**: `reference: '{case}'` adopts, for each case, +the grid of that case's *own* entry in `reference_group`. That is the registration idiom — +`reference: '{case}', reference_group: DVF` lands every moved image on its own field's grid, +which is where a displacement field is defined. A literal reference stays one header lookup +for the whole cohort; a per-case one is one per case, headers only either way. + + 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. A `reference` grid is what makes diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md index 83517586..b9d56797 100644 --- a/docs/source/reference/components/transforms.md +++ b/docs/source/reference/components/transforms.md @@ -128,6 +128,7 @@ Operate on a stacked `[N, …]` ensemble axis (prediction post-processing). | --- | --- | --- | --- | --- | | `InferenceStack` | Aggregate an ensemble stack (mean / median / seg-argmax); writes an `InferenceStack` volume. | `dataset, name, mode="mean"` | no | no — it writes the whole per-member stack | | `Norm` | Vector magnitude over the trailing axis (drops it). | — | **yes** | no‡ | +| `Magnitude` | Vector magnitude over the **channel** axis (`[C, …]` → `[1, …]`) — the channel-first sibling of `Norm`, for a stored vector volume such as a displacement field read as a case. | — | no† | **yes** — pointwise | | `Variance` | Per-voxel variance over N. | — | no† | **yes** | | `StandardDeviation` | Per-voxel std over N. | — | no† | **yes** | | `SegmentationDisagreement` | Per-voxel label disagreement across N segmentations. | `ignore_background=False` | no† | **yes** | @@ -139,8 +140,8 @@ Operate on a stacked `[N, …]` ensemble axis (prediction post-processing). | --- | --- | --- | | `Statistics` | Records ImageMin/Max/Mean/Std to the attribute cache and returns the tensor unchanged (feeds the perceptual criteria `SAM_Perceptual`, `IMPACTSynth`, `IMPACTReg`). Order in the transform list matters. | no‡ | | `Save` | Writes the preprocessed volume to a cache dataset and passes the tensor through. Once that cache exists it is read instead, and the transforms before it are skipped. A group written this way is also readable, within the same run, by anything that names it (`Mask: {path: }`), including when the write comes from a loader worker — every backend. On `h5` the reader and the writer share one store, which HDF5 does not define for concurrent access without SWMR: the entry is seen, but a read racing a write can raise. One file per case (`mha`/`nii`/…) has no such window. | no — it needs the whole volume to write | -| `Write` | A `Save` that is a **deliverable**: same boundary semantics, but `dataset` has no default, so a bare `Write:` fails at config time instead of writing into the source tree. The TRANSFORM workflow plans, resumes and reports on `Write` stages and requires every chain to end with one; a `Save` between them is an opportunistic milestone. Args: `dataset` (required), `group=None`, `scale_factors=None`, `downsample_method=None`. | region-writes where the backend allows (uncompressed `mha` / `h5` / `omezarr`) | -| `Reduce` | Folds every case of a group into one volume at fixed voxel — the stage that makes a chain N-to-1. `operator` is a classpath resolved against `konfai.data.reduction` (`Mean`, `Median`, `Vote` for label maps, `Concat`, or your own `Reduction`); `output` is required. `grid` is `strict` / `shape_only` / `reference:`. Args: `operator="Median", output="", grid="strict", grid_tolerance=1e-6, provenance=True`. **TRANSFORM only** — applied to one case it raises. | driven by the reduction engine, one region at a time | +| `Write` | A `Save` that is a **deliverable**: same boundary semantics, but `dataset` has no default, so a bare `Write:` fails at config time instead of writing into the source tree. The TRANSFORM workflow plans, resumes and reports on `Write` stages and requires every chain to end with one; a `Save` between them is an opportunistic milestone. Args: `dataset` (required), `group=None`, `scale_factors=None`, `downsample_method=None`. | region-writes where the backend allows (uncompressed `mha`/`nii`, `h5`, `omezarr`; `itktransform` for a displacement field) | +| `Reduce` | Folds every case of a group into one volume at fixed voxel — the stage that makes a chain N-to-1. `operator` is a classpath resolved against `konfai.data.reduction` (`Mean`, `Median`, `Std` — the ensemble-spread map, folded with running moments — `Vote` for label maps, `Concat`, or your own `Reduction`); `output` is required. `grid` is `strict` / `shape_only` / `reference:`. Args: `operator="Median", output="", grid="strict", grid_tolerance=1e-6, provenance=True`. **TRANSFORM only** — applied to one case it raises. | driven by the reduction engine, one region at a time | | `Expand` | Turns one case into `nb` copies at a declared point of the chain — `Reduce`'s mirror (1-to-N). Stages before it run once per case, stages after it once per copy, and the draws after it are ordinary stages. `pattern` is a `str.format` template and **both** `{name}` and `{a}` are required. Args: `nb=2, pattern="{name}_{a:02d}", seed=None`. **TRANSFORM only** — applied to one case it raises. | per copy, sharing one read pass when every per-copy stage is pointwise | | `KonfAIInference` | Run a nested KonfAI app inference in a spawned subprocess. Needs `konfai-apps` and `num_workers: 0`; defaults to a specific HF repo. | no‡ | diff --git a/docs/source/usage/index.rst b/docs/source/usage/index.rst index 3a75c425..9e89cd86 100644 --- a/docs/source/usage/index.rst +++ b/docs/source/usage/index.rst @@ -20,6 +20,7 @@ navigation: :doc:`apps` and :doc:`mcp`. :maxdepth: 1 adopting-konfai + python-workflows large-images custom-models docker diff --git a/docs/source/usage/python-workflows.md b/docs/source/usage/python-workflows.md new file mode 100644 index 00000000..190df151 --- /dev/null +++ b/docs/source/usage/python-workflows.md @@ -0,0 +1,60 @@ +# KonfAI in Python + +The four CLI commands are callables: `konfai.transform` (with `konfai.plan_transform`, its +dry-run twin), `konfai.evaluate`, `konfai.predict` and `konfai.train`. One engine, two spellings — +everything below builds the same config tree the YAML file would hold and hands it to the same +binder, so nothing here can drift from what a YAML run does. + +```python +import konfai +from konfai.data.transform import Resample, Write + +result = konfai.transform( + "moved", + "./Staged:mha", + {"Moving": {"Moved": [ + Resample(reference="{case}", reference_group="DVF", field_group="DVF"), + Write(dataset="./Output:mha"), + ]}}, + memory_budget="8G", +) +result.outputs # every chain's terminal Write: where the deliverables landed +result.config # the resolved YAML the run kept -- commit this file to version the experiment +``` + +A chain is a list of **live stage objects** — the very classes the YAML names, with the very same +constructor arguments, which the extension bases record as given — or the equivalent mapping +(`{"Resample": {...}, "Write": {...}}`), or a whole tree loaded from an existing YAML and modified +in place. Two stages of the same class in one chain spell the second one module-qualified +(`konfai.data.transform:Resample`), exactly as the YAML file must. + +## The contract, and how it differs from the CLI + +- **A designed refusal raises** `KonfAIError` — the message and the remedy are the exception; the + caller decides. Only the CLI catches and exits. +- **Results come back structured**: `transform` returns the `outputs.json` destinations and the + workspace; `evaluate` returns the parsed `Metric_*.json` as a dict. +- **The process is left as found**: the `KONFAI_*` environment is restored around every call, and + one workflow runs at a time per process — a second concurrent call is refused with the remedy + (subprocesses), never allowed to corrupt the first. +- **The record remains.** Every call materializes the resolved YAML in the run's workspace: + promoting a notebook run to a versioned experiment is copying `result.config` — nothing to + rewrite, and the run stays resumable like any other. + +`konfai.plan_transform(...)` takes the same arguments and returns the `TransformPlan` without +running anything — plan first is the same reflex in Python as on the CLI. + +## Which spelling fits which workflow + +| Workflow | Its config is… | The Python spelling | +| --- | --- | --- | +| TRANSFORM | a chain of stage objects | `konfai.transform(name, datasets, chains, ...)` with live stages | +| EVALUATION | criteria per group | `konfai.evaluate(name, datasets, metrics={"PRED": {"GT": [MAE(), Dice()]}}, ...)` | +| PREDICTION | wiring (checkpoints, patches, TTA) | `konfai.predict(models=[...], config=tree_or_path, ...)` — the tree or the file | +| TRAIN / RESUME | the full graph (model, losses, optimizers) | the tree: load the YAML into a dict, change the keys under study, call `konfai.train(config=tree)` | + +Every workflow entry point accepts the config **tree as a dict** wherever it accepts a file path — +that alone is the sweep idiom for TRAIN: the resolved config each run keeps *is* the record of what +was tried. The object spelling exists where a config is a list of objects (TRANSFORM chains, +EVALUATION criteria); rebuilding a training graph in nested kwargs would add nothing over the YAML +that publishes it. From 995fc4b3fdeebd4ca8d533153c985531c71d09a6 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 12:57:47 +0200 Subject: [PATCH 34/39] feat(impact-reg): the preset names its output, and nothing else is derived A preset declares one output: its transform, under whatever name and in whatever form its consumer reads. register discovers that group instead of assuming DVF, refuses an ensemble whose members disagree on it, and no longer derives Transform.h5 -- a preset that serves Slicer writes the ITK file itself, so there is no second copy of the same field under another name. Only the moved image is still derived. --- apps/impact_reg/README.md | 4 +- apps/impact_reg/impact_reg_konfai/cli.py | 4 +- .../impact_reg_konfai/impact_reg.py | 111 +++++++++++++----- .../tests/unit/test_displacement_field_io.py | 7 +- .../tests/unit/test_orchestration.py | 94 +++++++++++++-- .../tests/unit/test_tmp_dir_forwarding.py | 5 +- docs/source/usage/apps.md | 4 +- 7 files changed, 173 insertions(+), 56 deletions(-) diff --git a/apps/impact_reg/README.md b/apps/impact_reg/README.md index 6c3d8c1f..e9aeea55 100644 --- a/apps/impact_reg/README.md +++ b/apps/impact_reg/README.md @@ -84,7 +84,7 @@ The CLI is organised into sub-commands, matching the registration workflow: | Sub-command | Purpose | |---|---| -| `register` | Register a moving image onto a fixed image with one or more presets. Several presets are ensembled (their displacement fields are averaged). Writes the moved image, the displacement field (`DVF`), the transform, and the per-preset fields (kept for `uncertainty`). | +| `register` | Register a moving image onto a fixed image with one or more presets. Several presets are ensembled (their displacement fields are averaged). Writes the transform under the name and in the form the preset declared, the moved image derived from it, and — with `--keep_dvf` — the per-preset fields (kept for `uncertainty`). | | `eval` | Evaluate a registration on any subset of modalities — image (MAE), segmentation (Dice), landmarks (TRE). At least one modality is required. | | `uncertainty` | Voxel-wise spread map from an ensemble of displacement fields. | @@ -98,7 +98,7 @@ Evaluate a registration — any subset of modalities; the transform comes from a ```bash impact-reg-konfai eval \ - --transform ./Output/P000/Transform.h5 \ + --transform ./Output/P000/DVF.mha \ -f fixed.nii.gz -m moving.nii.gz --mask roi.nii.gz \ --gt-fixed-seg fixed_seg.nii.gz --gt-moving-seg moving_seg.nii.gz \ --gt-fixed-fid fixed.fcsv --gt-moving-fid moving.fcsv \ diff --git a/apps/impact_reg/impact_reg_konfai/cli.py b/apps/impact_reg/impact_reg_konfai/cli.py index 24d0b1f6..214423dd 100644 --- a/apps/impact_reg/impact_reg_konfai/cli.py +++ b/apps/impact_reg/impact_reg_konfai/cli.py @@ -132,8 +132,8 @@ def main() -> None: "--fields_only", dest="fields_only", action="store_true", - help="Write the displacement fields only: skip the moved image and Transform.h5, both derived " - "from the field. For a caller that composes the field itself and would delete them.", + help="Write the transforms only: skip the moved image, which is derived from them. For a " + "caller that composes the transform itself and would delete it.", ) _add_device(reg) _add_tmp_dir(reg) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 57c71b27..5060e2a6 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -58,6 +58,18 @@ _FORMATS = {".ome.zarr": "omezarr", ".zarr": "omezarr", ".h5": "itktransform", ".tfm": "itktransform"} +def _is_transform_file(path: Path) -> bool: + """An ITK transform file, as opposed to a displacement field stored as an image or a store.""" + return path.suffix in {".h5", ".tfm"} or path.name.endswith(".itk.txt") + + +def _as_transform(path: Path) -> "sitk.Transform": + """The stored registration as one ``sitk.Transform``, whatever form the preset wrote it in.""" + if _is_transform_file(path): + return sitk.ReadTransform(str(path)) + return sitk.DisplacementFieldTransform(sitk.Cast(read_displacement_field(path), sitk.sitkVectorFloat64)) + + def _app_id(preset: str) -> str: """Resolve a preset to a KonfAIApp id: a local ``/`` path, or ``:`` on HF.""" if Path(IMPACT_REG_KONFAI_REPO).is_dir(): @@ -86,6 +98,26 @@ def get_available_presets(force_update: bool = False) -> list[str]: return list(get_available_apps_on_hf_repo(IMPACT_REG_KONFAI_REPO, force_update)) +def _find_output_group(root: Path) -> str: + """The name of the single output group a preset produced under ``root``. + + A preset declares ONE output — its transform, in whatever form and under whatever name it chose. + konfai writes one dataset per output group (``///.``), so the group + is the directory holding the cases. Discovering it rather than assuming ``DVF`` is what lets an + official preset name its output ``Transform`` — where Slicer looks for it — while this pipeline's + own name theirs ``DVF``, with no branch here. + """ + runs = [child for child in sorted(root.iterdir()) if child.is_dir()] if root.is_dir() else [] + groups = [group for run in runs for group in sorted(run.iterdir()) if group.is_dir()] + if len(groups) != 1: + found = ", ".join(group.name for group in groups) or "none" + raise FileNotFoundError( + f"Expected the preset to produce exactly one output group under {root}, found {found}." + " A registration preset declares one output: its transform." + ) + return groups[0].name + + def _find_outputs(root: Path, stem: str) -> dict[str, Path]: """Every output named ``stem`` under ``root``, keyed by the CASE it belongs to. @@ -282,8 +314,8 @@ def _infer_preset( quiet: bool, tta: int = 0, config_overrides: list[str] | None = None, - ) -> dict[str, Path]: - """Run one preset app on every case at once; return its displacement field per case. + ) -> tuple[str, dict[str, Path]]: + """Run one preset app on every case at once; return its output group and its transform per case. ONE RUN, NOT ONE PER CASE. Each ``-i`` is an input GROUP, and konfai-apps expands each group's paths into units -- a file is one, a store or DICOM series is one, a plain directory is walked @@ -334,7 +366,8 @@ def _infer_preset( # be computed from it -- the moved image above all -- is this layer's job. Looking for a Moved # here would make every preset carry an output it does not owe, and a tiled one blend it across # every patch seam for a caller that has the field. - return _find_outputs(out, "DVF") + group = _find_output_group(out) + return group, _find_outputs(out, group) def register( self, @@ -364,11 +397,11 @@ def register( ``tmp_dir`` names where the intermediates are staged; see :func:`_work_dir`. - ``fields_only`` writes the displacement fields and stops there. The moved image and - ``Transform.h5`` are both derived FROM the field, at the cost of a full-size resample and a - full-size rewrite -- worth it for a caller that wants a registration to look at, waste for one - that composes the field with another and derives its own. A caller that reads only the fields - should be able to say so rather than pay for outputs it deletes. + ``fields_only`` writes the transforms and stops there. The moved image is derived FROM the + transform, at the cost of a full-size resample and a full-size rewrite -- worth it for a + caller that wants a registration to look at, waste for one that composes the field with + another and derives its own. A caller that reads only the fields should be able to say so + rather than pay for an output it deletes. """ # The cases are konfai-apps' to define, not ours to count. It expands each input GROUP into # units -- a file, a store, a DICOM series, or every volume inside a plain directory -- and pairs @@ -403,6 +436,14 @@ def register( ) for preset in presets } + groups = {group for group, _ in fields_by_preset.values()} + if len(groups) != 1: + raise RuntimeError( + f"the presets named their output differently ({', '.join(sorted(groups))}); an" + " ensemble folds one group, so every member must declare the same one." + ) + group = groups.pop() + fields_by_preset = {preset: fields for preset, (_, fields) in fields_by_preset.items()} cases = sorted(fields_by_preset[presets[0]]) for preset, fields in fields_by_preset.items(): if sorted(fields) != cases: @@ -434,27 +475,26 @@ def register( dvf_paths.append(dvf) if len(presets) == 1: - _copy_output(dvf_paths[0], case_out, "DVF") + _copy_output(dvf_paths[0], case_out, group) else: # Ensemble: fold the presets' fields (all on the fixed grid -- and Reduce VERIFIES # the claim) into the averaged DVF, the one output no single preset produced. # Streamed: the fold is incremental, so the peak is one accumulator plus the # member being read, whatever the size of the ensemble. - self._ensemble_mean(case, presets, dvf_paths, output, work, gpu, cpu, quiet) + self._ensemble_mean(case, group, presets, dvf_paths, output, work, gpu, cpu, quiet) if not fields_only: - # The moved images AND Transform.h5 in ONE streamed run: Resample adopts, per case, - # the grid of that case's own DVF (a field is defined ON the fixed grid) and reads - # the field as the map; the second chain writes the same field as an ITK transform - # file (the ':itktransform' backend fills it region by region) -- one plan, both - # deliverables, resumable. - self._derive_moved(dict(zip(cases, moving_units, strict=True)), output, work, gpu, cpu, quiet) + # The moved images in ONE streamed run over the cohort: Resample adopts, per case, + # the grid of that case's own field (a field is defined ON the fixed grid) and reads + # the field as the map. + self._derive_moved(dict(zip(cases, moving_units, strict=True)), group, output, work, gpu, cpu, quiet) finally: shutil.rmtree(work, ignore_errors=True) def _ensemble_mean( self, case: str, + group: str, presets: list[str], dvf_paths: list[Path], output: Path, @@ -463,20 +503,20 @@ def _ensemble_mean( cpu: int | None, quiet: bool, ) -> None: - """Average one case's preset fields into ``//DVF`` — Reduce(Mean), streamed.""" + """Average one case's preset fields into ``//`` — Reduce(Mean), streamed.""" from konfai.data.transform import Reduce, Write members = _stage_group( work / f"ensemble_{case}", "DVF", {preset: dvf for preset, dvf in zip(presets, dvf_paths, strict=True)} ) suffixes = "".join(dvf_paths[0].suffixes) - _output_path(output / case, "DVF", suffixes) # drop a stale other-form DVF before writing + _output_path(output / case, group, suffixes) # drop a stale other-form output before writing _run_transform( f"impact_reg_ensemble_{case}", [members], { "DVF": { - "DVF": [ + group: [ Reduce(operator="Mean", output=case, grid="strict"), Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), ] @@ -491,27 +531,32 @@ def _ensemble_mean( def _derive_moved( self, cases: dict[str, Path], + group: str, output: Path, work: Path, gpu: list[int], cpu: int | None, quiet: bool, ) -> None: - """The moved images and ``Transform.h5``, both derived from the fields — one run, two chains. + """The moved images, derived from the transforms — one streamed run over the whole cohort. A preset that emits only a field is complete: everything else IS that field. The moved image: ``reference: '{case}'`` adopts each case's own DVF grid (a field is defined ON the fixed grid) and the field is the map (``field_group``) — one interpolation, streamed, each - slab's source window sized from the field values it reads. ``Transform.h5``: the same field - written as an ITK transform file, a plain ``Write`` the ``:itktransform`` backend fills - region by region. + slab's source window sized from the field values it reads. + + NOTHING ELSE IS DERIVED. A preset writes its transform in the form its consumer reads — an ITK + transform file where Slicer picks it up, an RFC-5 store where this pipeline streams it — so + there is no second copy of the same field to produce under another name. """ from konfai.data.transform import Resample, Write - fields = {case: _the_output(output / case, "DVF") for case in cases} - suffixes = "" - for case, dvf in fields.items(): - suffixes = "".join(dvf.suffixes) + fields = {case: _the_output(output / case, group) for case in cases} + # The moved image is written in the MOVING's own form: the fields' form may be an ITK + # transform file, which cannot hold an image. _stage_group refuses a mixed Moving group + # below, so the first case's form speaks for the cohort. + suffixes = "".join(next(iter(cases.values())).suffixes) + for case in fields: _output_path(output / case, "Moved", suffixes) # drop a stale other-form Moved moving_root = _stage_group(work / "moved_stage", "Moving", cases) field_root = _stage_group(work / "moved_stage", "DVF", fields) @@ -525,7 +570,6 @@ def _derive_moved( Write(dataset=f"{output}:{_FORMATS.get(suffixes.lower(), suffixes.lstrip('.'))}"), ] }, - "DVF": {"Transform": [Write(dataset=f"{output}:itktransform")]}, }, work, gpu, @@ -614,8 +658,7 @@ def evaluate( if index < len(gt_fixed_fid) and index < len(gt_moving_fid): fixed_points = read_landmarks(gt_fixed_fid[index]) if transform_path is not None: - transform = sitk.ReadTransform(str(transform_path)) - fixed_points = apply_to_data_transform(fixed_points, {transform: False}) + fixed_points = apply_to_data_transform(fixed_points, {_as_transform(transform_path): False}) moved_fid = work / "moved_fid.fcsv" write_landmarks(fixed_points, moved_fid) app.evaluate( @@ -661,7 +704,10 @@ def _warp_onto_fixed( resample["interpolation"] = "nearest" if transform_path is not None: datasets.append(_stage_group(base, "Reg", {"P000": transform_path})) - resample["transforms"] = {"Reg": False} + if _is_transform_file(transform_path): + resample["transforms"] = {"Reg": False} + else: + resample["field_group"] = "Reg" out_root = work / f"moved_{kind}" _run_transform( f"impact_reg_eval_{kind}", @@ -702,7 +748,8 @@ def uncertainty( members = _units(list(dvfs)) spec = _stage_group(work / "members", "DVF", {f"M{index:03d}": dvf for index, dvf in enumerate(members)}) - suffixes = "".join(members[0].suffixes) + suffixes = ".mha" if _is_transform_file(members[0]) else "".join(members[0].suffixes) + _output_path(output / "uncertainty", "Uncertainty", suffixes) # drop a stale other-form map _run_transform( "impact_reg_uncertainty", [spec], diff --git a/apps/impact_reg/tests/unit/test_displacement_field_io.py b/apps/impact_reg/tests/unit/test_displacement_field_io.py index d8efdb30..fdc92f07 100644 --- a/apps/impact_reg/tests/unit/test_displacement_field_io.py +++ b/apps/impact_reg/tests/unit/test_displacement_field_io.py @@ -223,9 +223,8 @@ def test_rerunning_in_the_other_form_leaves_one_output(tmp_path: Path) -> None: def test_ensemble_field_written_by_the_orchestrator_is_a_declared_field(tmp_path: Path) -> None: """The averaged DVF is folded by Reduce(Mean) and written through konfai's Write, in the members' form — and must stay a DECLARED field: ``read_displacement_field`` refuses an - undeclared 3-channel store, so dropping the declaration would break Transform.h5 and - ``evaluate`` right after a successful register. The values are the voxel-wise mean, on the - members' geometry.""" + undeclared 3-channel store, so dropping the declaration would break ``evaluate`` right after + a successful register. The values are the voxel-wise mean, on the members' geometry.""" members = [] for index in (1, 2): (tmp_path / f"m{index}").mkdir() @@ -233,7 +232,7 @@ def test_ensemble_field_written_by_the_orchestrator_is_a_declared_field(tmp_path output, work = tmp_path / "out", tmp_path / "work" work.mkdir() - ImpactRegKonfAIApp()._ensemble_mean("P000", ["a", "b"], members, output, work, [], 1, True) + ImpactRegKonfAIApp()._ensemble_mean("P000", "DVF", ["a", "b"], members, output, work, [], 1, True) out = output / "P000" / "DVF.ome.zarr" assert is_displacement_field(out) diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index 8343cfad..d8e59452 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -92,7 +92,7 @@ def test_ensemble_mean_is_the_voxelwise_mean_with_reference_geometry(tmp_path: P output, work = tmp_path / "out", tmp_path / "work" work.mkdir() - reg.ImpactRegKonfAIApp()._ensemble_mean("P000", ["a", "b"], paths, output, work, [], 1, True) + reg.ImpactRegKonfAIApp()._ensemble_mean("P000", "DVF", ["a", "b"], paths, output, work, [], 1, True) avg = sitk.ReadImage(str(output / "P000" / "DVF.mha")) field = sitk.GetArrayFromImage(avg) @@ -105,14 +105,15 @@ def test_ensemble_mean_is_the_voxelwise_mean_with_reference_geometry(tmp_path: P def _stub_infer(app: reg.ImpactRegKonfAIApp, moving_image: Path, dvf_by_preset: dict[str, tuple]): - """Replace ``_infer_preset`` so it writes a constant DVF.mha per preset on the moving grid.""" + """Replace ``_infer_preset`` so it writes a constant DVF.mha per preset on the moving grid, + reported under group ``DVF`` as the real one reports the group it discovered.""" reference = sitk.ReadImage(str(moving_image)) def fake(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): out = Path(work) / preset / "P000" out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", dvf_by_preset[preset], reference) - return {"P000": out / "DVF.mha"} + return "DVF", {"P000": out / "DVF.mha"} app._infer_preset = fake # type: ignore[method-assign] @@ -129,7 +130,9 @@ def test_register_single_preset_reuses_the_field_and_derives_the_moved(tmp_path: app.register(["FireANTs_SyN"], [fixed], [moving], output=out) case = out / "P000" - assert (case / "Moved.mha").is_file() and (case / "DVF.mha").is_file() and (case / "Transform.h5").is_file() + assert (case / "Moved.mha").is_file() and (case / "DVF.mha").is_file() + # nothing else is derived: the transform exists only in the form the preset wrote it + assert not (case / "Transform.h5").exists() # single preset: the DVF is the model's own field, reused verbatim (no re-averaging) field = sitk.GetArrayFromImage(sitk.ReadImage(str(case / "DVF.mha"))) np.testing.assert_allclose(field[0, 0, 0], (2.0, 0.0, 0.0), atol=1e-6) @@ -174,7 +177,7 @@ def field_only(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, out = Path(work) / preset / "P000" out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) - return {"P000": out / "DVF.mha"} + return "DVF", {"P000": out / "DVF.mha"} app._infer_preset = field_only # type: ignore[method-assign] out = tmp_path / "Output" @@ -191,9 +194,9 @@ def field_only(preset, fixed, moving, fixed_masks, moving_masks, n_cases, work, def test_register_fields_only_writes_nothing_derived(tmp_path: Path) -> None: """A caller that composes the field itself pays for the field, and nothing else. - Both the moved image and Transform.h5 are derived FROM the field -- a full-size resample and a - full-size rewrite of the same voxels. The tiled refinement reads the field, composes it with its - global pass and derives its own moved, so producing them for it is pure waste. + The moved image is derived FROM the field -- a full-size resample of the same voxels. The tiled + refinement reads the field, composes it with its global pass and derives its own moved, so + producing one for it is pure waste. """ moving = tmp_path / "moving.mha" sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) @@ -232,17 +235,84 @@ def field_only(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, wo out = Path(work) / preset / "P000" out.mkdir(parents=True, exist_ok=True) _write_dvf(out / "DVF.mha", (2.0, 0.0, 0.0), reference) - return {"P000": out / "DVF.mha"} + return "DVF", {"P000": out / "DVF.mha"} app._infer_preset = field_only # type: ignore[method-assign] out = tmp_path / "Output" app.register(["FireANTs_SyN"], [fixed], [moving], output=out) - # moved(p) = moving(p + d), d = +2 along x on a unit grid: moving is z*64 + y*8 + x. - moved = sitk.GetArrayFromImage(sitk.ReadImage(str(out / "P000" / "Moved.mha"))) + # The moved image takes the MOVING's form -- a store in, a store out -- while the field the + # preset wrote keeps its own. moved(p) = moving(p + d), d = +2 along x on a unit grid: + # moving is z*64 + y*8 + x. + store = out / "P000" / "Moved.ome.zarr" + assert store.is_dir(), "the moved image was not written in the moving's own form" + moved = ome_zarr.read_ome_zarr_data_slice(store, (slice(None),) * 4)[0][0] np.testing.assert_allclose(moved[0, 0, 0], 2.0, atol=1e-6) np.testing.assert_allclose(moved[1, 1, 0], 74.0, atol=1e-6) - assert (out / "P000" / "Transform.h5").is_file() + assert (out / "P000" / "DVF.mha").is_file() + + +def test_register_adopts_the_presets_output_name(tmp_path: Path) -> None: + """A preset names its output; the pipeline follows. An official preset calls its transform + ``Transform`` — where Slicer looks for it — and ``register`` must not rename it ``DVF``.""" + moving = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + reference = sitk.ReadImage(str(moving)) + app = reg.ImpactRegKonfAIApp() + + def named_transform(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + out = Path(work) / preset / "P000" + out.mkdir(parents=True, exist_ok=True) + _write_dvf(out / "Transform.mha", (2.0, 0.0, 0.0), reference) + return "Transform", {"P000": out / "Transform.mha"} + + app._infer_preset = named_transform # type: ignore[method-assign] + out = tmp_path / "Output" + app.register(["FireANTs_SyN"], [fixed], [moving], output=out) + + case = out / "P000" + assert (case / "Transform.mha").is_file() and (case / "Moved.mha").is_file() + assert not (case / "DVF.mha").exists() + + +def test_register_refuses_presets_that_name_their_output_differently(tmp_path: Path) -> None: + """An ensemble folds one group: members that disagree on its name are refused, not renamed.""" + moving = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + reference = sitk.ReadImage(str(moving)) + app = reg.ImpactRegKonfAIApp() + group_by_preset = {"A": "DVF", "B": "Transform"} + + def mixed(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + group = group_by_preset[preset] + out = Path(work) / preset / "P000" + out.mkdir(parents=True, exist_ok=True) + _write_dvf(out / f"{group}.mha", (2.0, 0.0, 0.0), reference) + return group, {"P000": out / f"{group}.mha"} + + app._infer_preset = mixed # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match="named their output differently"): + app.register(["A", "B"], [fixed], [moving], output=tmp_path / "Output") + + +def test_find_output_group_discovers_the_one_group(tmp_path: Path) -> None: + """konfai-apps writes ``///…``; the group is the one directory holding cases.""" + (tmp_path / "reg" / "Transform" / "P000").mkdir(parents=True) + assert reg._find_output_group(tmp_path) == "Transform" + + +def test_find_output_group_refuses_more_than_one(tmp_path: Path) -> None: + (tmp_path / "reg" / "DVF" / "P000").mkdir(parents=True) + (tmp_path / "reg" / "Moved" / "P000").mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="exactly one output group"): + reg._find_output_group(tmp_path) def test_stage_group_replaces_an_existing_link(tmp_path: Path) -> None: diff --git a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py index 7b244b6a..f99f2611 100644 --- a/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py +++ b/apps/impact_reg/tests/unit/test_tmp_dir_forwarding.py @@ -48,8 +48,9 @@ def test_infer_preset_forwards_the_workspace(tmp_path: Path, monkeypatch, write_ def fake_run(command, **kwargs): captured.append(list(command)) - # Stand in for the preset run: konfai-apps leaves one case directory per unit under -o. - write_preset_output(Path(command[command.index("-o") + 1]) / "P000") + # Stand in for the preset run: konfai-apps leaves one dataset per output group under -o, + # laid out // -- the shape _find_output_group discovers the group from. + write_preset_output(Path(command[command.index("-o") + 1]) / "reg" / "DVF" / "P000") return None monkeypatch.setattr("impact_reg_konfai.impact_reg.subprocess.run", fake_run) diff --git a/docs/source/usage/apps.md b/docs/source/usage/apps.md index 4d354db8..1a9467da 100644 --- a/docs/source/usage/apps.md +++ b/docs/source/usage/apps.md @@ -261,8 +261,8 @@ voxel values remain exactly identical, verified one plane at a time. This creates a known, reproducible offset without inventing anatomy or interpolating an input image. -The App writes `Moved.mha`, a three-component `DVF.mha` in millimetres, and a -reusable `Transform.h5` on the fixed CT grid. Compared with the unshifted MR, +The App writes the transform the preset declared — here a three-component +`DVF.mha` in millimetres — and `Moved.mha`, both on the fixed CT grid. Compared with the unshifted MR, foreground NCC improves from `0.129` to `0.937` and MAE from `106.11` to `21.09`. The field has a mean magnitude of `23.06 mm` and a 95th percentile of `25.55 mm`. These validation values were accumulated one axial plane at a time. From 9b93e0c275a40c9adc26400aaf07b7ae27524482 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 15:19:27 +0200 Subject: [PATCH 35/39] fix(data)!: a backend token is not an extension, and h5py is not optional `:itktransform` writes `.h5` -- nothing on disk is ever named `.itktransform`. Validating a dataset spec against the extension list therefore rejected the very format the write side had just produced: a run could write a transform it could not read back. Tokens that name a backend rather than a suffix now live in SUPPORTED_BACKEND_FORMATS, and SUPPORTED_FORMATS -- the union -- is what a `path[:flag]:format` spec is checked against, down to `split_path_spec`'s parameter. SUPPORTED_EXTENSIONS keeps its own job: the suffixes probed beside a case and matched against an input path. The backend also stops degrading in silence. Without h5py it fell back to holding the field whole in float64 and writing it through sitk, so peak memory turned on whether an optional import had succeeded -- and neither path was covered, the guarantee test being skipped when h5py is absent. It now raises and names the extra, which is what the `h5` backend has always done in effect (an AttributeError on None; it says so properly now too). impact-reg declares h5py: every preset writes its transform through this backend. --- AGENTS.md | 2 +- apps/impact_reg/setup.py | 4 +- docs/source/concepts/datasets.md | 5 ++- .../reference/components/storage-backends.md | 13 ++++-- konfai-apps/konfai_apps/app.py | 4 +- konfai/data/data_manager.py | 10 ++--- konfai/data/patching.py | 4 +- konfai/utils/dataset.py | 41 +++++++------------ konfai/utils/utils.py | 20 +++++++-- tests/unit/test_imaging_formats.py | 8 ++-- tests/unit/test_itk_transform_backend.py | 34 +++++++++++++++ 11 files changed, 94 insertions(+), 51 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 27e3ef2d..9f07560b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -64,7 +64,7 @@ Every extension point is **"subclass a base, reference it by classpath in YAML"* - **Transform:** subclass `data.transform.Transform`; implement `__call__` **and** `transform_shape()` (must predict the output spatial shape *exactly* — patch planning depends on it). Declare `patch_locality()` (a `LocalityKind`: `POINTWISE`/`HALO`/`ORIENTATION`/`CROP`/`GLOBAL_STAT`/`REGRID`/`SLAB`/`WHOLE_VOLUME`) or the base default makes it `WHOLE_VOLUME`; a `WHOLE_VOLUME` that is a property of the *configuration* rather than of the stage must carry `reason=`, which the plan prints — without it the reader has nothing to change. Pair `inverse()` if `apply_inverse`; override `prepare(konfai_args)` only when the stage builds a sub-object from configuration of its own (`Reduce` → its operator). - **Augmentation:** subclass `data.augmentation.DataAugmentation`; `_state_init` (sample params per case index) + `_compute` (apply lazily). Only `Mask`/`Permute` may change shape. A draw is also a **chain stage**: `TransformLoader` resolves a bare name against `data.transform` first and `data.augmentation` second, so a `transforms:` block may interleave draws and transforms — which is how TRANSFORM declares per-copy draws after an `Expand`. - **Reduction:** subclass `data.reduction.Reduction`; implement `__call__(list[Tensor]) -> Tensor` over the `[1, K, C, *spatial]` layout both engines hand over. Two consumers, one vocabulary: the predictor folds one case's copies (ensemble/TTA), `data.transform.Reduce` folds N **cases** into one. Declare `voxel_local = True` only if every output voxel reads the same voxel of each input (**a wrong `True` corrupts a streamed output** — the gate checks nothing else), `incremental = True` if `accumulate` can fold one at a time, and override `output_channels(channels, cases)` when the fold changes the channel count (`Concat` does). `Reduce` refuses a non-`voxel_local` operator outright. -- **Imaging format:** add a `Dataset.AbstractFile` backend, dispatch it in `File.__enter__`, register aliases in `SUPPORTED_EXTENSIONS`; import-guard the heavy lib. +- **Imaging format:** add a `Dataset.AbstractFile` backend, dispatch it in `File.__enter__`, register aliases in `SUPPORTED_EXTENSIONS` — or in `SUPPORTED_BACKEND_FORMATS` when the token is not a suffix any file carries (`:itktransform` writes `.h5`), since only the extensions are probed on disk; import-guard the heavy lib and raise a `DatasetManagerError` naming the extra rather than degrading in silence. **Classpaths:** a bare name (e.g. `Dice`) resolves inside that kind's package; `module:Class` imports *any* module — a local file (`Loss:MyWrapper`) or an installed library (`monai.losses:DiceLoss`, `torch:nn:L1Loss`). Model classpaths resolve against `konfai.models.python`. The pre-1.6.0 absolute form `konfai.models..:` still resolves via a rewrite + `DeprecationWarning`; new code uses the relative or `default|` form. diff --git a/apps/impact_reg/setup.py b/apps/impact_reg/setup.py index fdcd68db..914e47f8 100644 --- a/apps/impact_reg/setup.py +++ b/apps/impact_reg/setup.py @@ -33,4 +33,6 @@ def _release_version() -> str: _version = _release_version() -setup(install_requires=[f"konfai=={_version}", f"konfai-apps=={_version}"]) +# h5py: every preset writes its transform through konfai's ':itktransform' backend, which fills the +# parameters region by region rather than holding the field in float64. The backend requires it. +setup(install_requires=[f"konfai=={_version}", f"konfai-apps=={_version}", "h5py"]) diff --git a/docs/source/concepts/datasets.md b/docs/source/concepts/datasets.md index 8c1430c3..c2545fa8 100644 --- a/docs/source/concepts/datasets.md +++ b/docs/source/concepts/datasets.md @@ -31,7 +31,10 @@ Dataset/ ``` The concrete file extension is not restricted to `.mha`. KonfAI supports the -extensions listed in `konfai.utils.utils.SUPPORTED_EXTENSIONS`. +extensions listed in `konfai.utils.utils.SUPPORTED_EXTENSIONS`. A spec may also +name a format that is a **backend rather than a suffix** (`:itktransform`, whose +entries are `.h5`): those live in `SUPPORTED_BACKEND_FORMATS`, and +`SUPPORTED_FORMATS` is the union a `path[:flag]:format` spec is checked against. Directory-backed formats use the same case/group model: diff --git a/docs/source/reference/components/storage-backends.md b/docs/source/reference/components/storage-backends.md index e07e6a22..c350ec79 100644 --- a/docs/source/reference/components/storage-backends.md +++ b/docs/source/reference/components/storage-backends.md @@ -22,7 +22,7 @@ whole). | `Dataset.H5File` | `h5` | Single monolithic HDF5 file | yes (chunked) | yes | `konfai[hdf5]` (`h5py`) | | `Dataset.OmeZarrFile` | `omezarr, ome-zarr, ome_zarr, zarr` (+ `@level`) | OME-Zarr pyramid directory | yes (chunked) | yes, `scale_factors` pyramids included | `konfai[omezarr]` (`zarr` + `ngff-zarr`) | | `Dataset.DicomFile` (DICOM series; scalar-array writes) | `dicom` | DICOM series directory | per slice | no (whole series) | `konfai[dicom]` (`pydicom`) | -| `Dataset.ItkTransformFile` | `itktransform` | ITK transform files (`.h5`, `.tfm`), one per case/group | yes for a displacement entry (the parameters are HDF5; a row span is one contiguous read) | yes — the parameters fill region by region, and the file is what `sitk.WriteTransform` would have written | `konfai[itk]`; `h5py` for the region paths (whole-file sitk fallback without it) | +| `Dataset.ItkTransformFile` | `itktransform` | ITK transform files (`.h5`, `.tfm`), one per case/group | yes for a displacement entry (the parameters are HDF5; a row span is one contiguous read) | yes, **for a 3-component 3-D displacement field with image geometry** — the parameters fill region by region, and the file is what `sitk.WriteTransform` would have written; any other transform kind is whole-entry | `konfai[itk]` + `konfai[hdf5]` (the parameters are touched through `h5py`) | ```{tip} `pip install "konfai[imaging]"` installs every backend at once @@ -40,9 +40,14 @@ write side is the point: a displacement field streams into the exact file — without ever holding the field whole in float64. The read side hands back what `Dataset.read_transform` decodes: a displacement entry as its field (region reads included), any other stored transform (affine, composite, `.tfm`) as its -parameters. This is how a run's `Transform.h5` deliverable is a plain `Write:` -like any other, and how a staged transform file resolves through the same -`Dataset` surface as an image. +parameters. This is how a registration preset's transform deliverable — under whatever +name it declares — is a plain `Write:` like any other, and how a staged +transform file resolves through the same `Dataset` surface as an image. + +`itktransform` is a **backend token, not an extension**: the file on disk is +`.h5` (or `.tfm`), and nothing is ever named `.itktransform`. That is the +one thing `SUPPORTED_BACKEND_FORMATS` exists to say — a format may be declared in +a dataset spec without being a suffix any path can carry. ## The `SitkFile` default backend also handles sidecars diff --git a/konfai-apps/konfai_apps/app.py b/konfai-apps/konfai_apps/app.py index 81f6df97..40656078 100644 --- a/konfai-apps/konfai_apps/app.py +++ b/konfai-apps/konfai_apps/app.py @@ -37,7 +37,7 @@ from konfai.utils.dataset import Dataset from konfai.utils.errors import AppRepositoryError, KonfAIAppClientError from konfai.utils.runtime import MinimalLog, State, safe_torch_load -from konfai.utils.utils import SUPPORTED_EXTENSIONS, split_format_level, split_path_spec +from konfai.utils.utils import SUPPORTED_EXTENSIONS, SUPPORTED_FORMATS, split_format_level, split_path_spec from ruamel.yaml import YAML from .app_repository import LocalAppRepository, get_app_repository_info @@ -1062,7 +1062,7 @@ def _dataset_level(prediction_file: str, dataset_dir: Path) -> int: str(entry), default_format="mha", allowed_flags={"a", "i"}, - supported_extensions=SUPPORTED_EXTENSIONS, + supported_formats=SUPPORTED_FORMATS, ) if Path(filename).resolve() == target: return split_format_level(file_format)[1] diff --git a/konfai/data/data_manager.py b/konfai/data/data_manager.py index 818aac2b..5a99ac8d 100755 --- a/konfai/data/data_manager.py +++ b/konfai/data/data_manager.py @@ -61,7 +61,7 @@ get_memory_info, memory_forecast, ) -from konfai.utils.utils import SUPPORTED_EXTENSIONS, OverlapSpec, resolve_patch, split_path_spec +from konfai.utils.utils import SUPPORTED_FORMATS, OverlapSpec, resolve_patch, split_path_spec # A cached case is a float32 tensor (torch's default dtype, and the default TensorCast's target), so # bytes are counted at 4/element from the header shape alone -- not the on-disk dtype, and without @@ -1251,14 +1251,14 @@ def _resolve_dataset_sources(self) -> dict[str, list[tuple[str, bool]]]: dataset_filename, default_format="mha", allowed_flags={"a", "i"}, - supported_extensions=SUPPORTED_EXTENSIONS, + supported_formats=SUPPORTED_FORMATS, ) append = flag != "i" - if file_format.split("@", 1)[0] not in SUPPORTED_EXTENSIONS: + if file_format.split("@", 1)[0] not in SUPPORTED_FORMATS: raise DatasetManagerError( f"Unsupported file format '{file_format}'.", - f"Supported extensions are: {', '.join(SUPPORTED_EXTENSIONS)}", + f"Supported formats are: {', '.join(SUPPORTED_FORMATS)}", ) dataset = Dataset(filename, file_format) @@ -1984,7 +1984,7 @@ def _output_destinations(self) -> dict[tuple[str, str], list[tuple[str, str]]]: for transform in group_transform.transforms: if isinstance(transform, Save) and transform.dataset: filename, _flag, _file_format = split_path_spec( - transform.dataset, default_format="mha", supported_extensions=SUPPORTED_EXTENSIONS + transform.dataset, default_format="mha", supported_formats=SUPPORTED_FORMATS ) entries.append((str(Path(filename).resolve()), transform.group or group_dest)) destinations[(group_src, group_dest)] = entries diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 73db817f..5217c051 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -45,7 +45,7 @@ from konfai.utils.dataset import Attribute, Dataset, DataStream from konfai.utils.errors import ConfigError, PatchError from konfai.utils.utils import ( - SUPPORTED_EXTENSIONS, + SUPPORTED_FORMATS, OverlapSpec, best_sweep_axis, concretize_patch_size, @@ -306,7 +306,7 @@ def save_destination(save: Save, default_dataset: Dataset, default_group: str) - filename, _, file_format = split_path_spec( save.dataset, default_format="mha", - supported_extensions=SUPPORTED_EXTENSIONS, + supported_formats=SUPPORTED_FORMATS, ) dataset = Dataset(filename, file_format, save.scale_factors, save.downsample_method) else: diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 5e321bac..a0c8378e 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -729,8 +729,6 @@ def _create_itk_transform_file(path: str, spatial: list[int], attributes: Attrib and the parameters, the field buffer with the component fastest, float64. Returns the open file and the parameters dataset. """ - import h5py - fixed = np.concatenate( [ np.asarray(spatial[::-1], dtype=np.float64), # size, in (x, y, z) @@ -1075,6 +1073,11 @@ class H5File(AbstractFile): _READ_CHUNK_CACHE_SLOTS = 100003 def __init__(self, filename: str, read: bool) -> None: + if h5py is None: + raise DatasetManagerError( + "An ':h5' dataset needs h5py.", + "Install it with: pip install konfai[hdf5]", + ) self.h5: h5py.File | None = None self.filename = filename if not self.filename.endswith(".h5"): @@ -1997,9 +2000,18 @@ class ItkTransformFile(AbstractFile): back what ``Dataset.read_transform`` decodes: a displacement entry carries its field and the displacement marker; any other stored transform, the parameter rows and type keys of ``_encode_transform_leaves``. + + Needs ``h5py``, as the ``h5`` backend does: the whole point is to touch the parameters + region by region, and a run whose peak memory turns on whether an optional import + succeeded is a run nobody can size. """ def __init__(self, filename: str, read: bool) -> None: + if h5py is None: + raise DatasetManagerError( + "An ':itktransform' dataset needs h5py.", + "Install it with: pip install konfai[hdf5]", + ) self.filename = filename self.read = read @@ -2032,10 +2044,6 @@ def file_to_data(self, group: str, name: str) -> tuple[np.ndarray, Attribute]: ) def bounded_region_reads(self, name: str) -> bool: - try: - import h5py # noqa: F401 - except ImportError: - return False shape, _attributes = self.get_infos("", name) return len(shape) == 4 and shape[0] == 3 @@ -2046,11 +2054,6 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) - is one contiguous span of the parameters: read, reshaped, and sliced down to the exact region — the peak is the row span, never the field. """ - try: - import h5py - except ImportError: - data, attributes = self.file_to_data(group, name) - return data[slices], attributes shape, attributes = self.get_infos(group, name) if len(shape) != 4 or shape[0] != 3 or DISPLACEMENT_FIELD_ATTRIBUTE not in attributes: data, attributes = self.file_to_data(group, name) @@ -2096,13 +2099,6 @@ def data_to_file( f" shape {list(array.shape)}.", "Write the field itself (channel-first, with its geometry), or a sitk.Transform.", ) - try: - import h5py # noqa: F401 - except ImportError: - field = sitk.Cast(data_to_image(array, attributes), sitk.sitkVectorFloat64) - sitk.WriteTransform(sitk.DisplacementFieldTransform(field), staging) - os.replace(staging, final) - return spatial = [int(extent) for extent in array.shape[1:]] file, parameters = _create_itk_transform_file(staging, spatial, attributes) with file: @@ -2118,10 +2114,6 @@ def open_data_stream( region_shape: list[int] | None = None, ) -> DataStream | None: del dtype, region_shape # the parameters are float64 whatever arrives, converted per slab - try: - import h5py # noqa: F401 - except ImportError: - return None if len(shape) != 4 or shape[0] != 3 or not is_an_image(attributes): return None os.makedirs(self.filename, exist_ok=True) @@ -2141,11 +2133,6 @@ def is_exist(self, group: str, name: str | None = None) -> bool: return os.path.exists(self._path(name if name else group)) def get_infos(self, group: str, name: str) -> tuple[list[int], Attribute]: - try: - import h5py - except ImportError: - data, attributes = self.file_to_data(group, name) - return [int(extent) for extent in data.shape], attributes with h5py.File(self._path(name), "r") as file: kind = bytes(file["TransformGroup/0/TransformType"][0]) fixed = np.asarray(file["TransformGroup/0/TransformFixedParameters"][()], dtype=np.float64) diff --git a/konfai/utils/utils.py b/konfai/utils/utils.py index e3c78f8e..8e6e8090 100755 --- a/konfai/utils/utils.py +++ b/konfai/utils/utils.py @@ -364,6 +364,8 @@ def get_patch_slices_from_shape( return _sweep_first(slices, sweep_axis), nb_patch_per_dim +# Suffixes an entry can carry on disk: probed next to a case to find an entry whatever it was +# written as, and matched against a path to recognise an input file. SUPPORTED_EXTENSIONS = [ "mha", "mhd", # MetaImage @@ -395,6 +397,16 @@ def get_patch_slices_from_shape( "npy", ] +# Format tokens that name a backend rather than a suffix. An ':itktransform' entry is one ITK +# transform file, written as `.h5`, so nothing on disk ever ends in `.itktransform`: the +# token is legal wherever a format is declared, and must never be probed as an extension. +SUPPORTED_BACKEND_FORMATS = [ + "itktransform", +] + +# Everything a `path[:flag]:format` spec may name. +SUPPORTED_FORMATS = [*SUPPORTED_EXTENSIONS, *SUPPORTED_BACKEND_FORMATS] + _WINDOWS_ABSOLUTE_PATH_RE = re.compile(r"^[A-Za-z]:[\\/]") @@ -423,7 +435,7 @@ def split_path_spec( *, default_format: str = "mha", allowed_flags: set[str] | None = None, - supported_extensions: list[str] | None = None, + supported_formats: list[str] | None = None, ) -> tuple[str, str | None, str]: """Split a KonfAI ``path[:flag]:format`` spec without breaking Windows paths. @@ -439,7 +451,7 @@ def split_path_spec( is preserved. """ - extensions = SUPPORTED_EXTENSIONS if supported_extensions is None else supported_extensions + formats = SUPPORTED_FORMATS if supported_formats is None else supported_formats parts = value.rsplit(":", 2) if len(parts) == 1: @@ -447,14 +459,14 @@ def split_path_spec( if len(parts) == 2: path, maybe_format = parts - if maybe_format in extensions: + if maybe_format in formats: return path, None, maybe_format if is_windows_absolute_path(value): return value, None, default_format return path, None, maybe_format path, middle, file_format = parts - if file_format in extensions: + if file_format in formats: if allowed_flags is not None and middle in allowed_flags: return path, middle, file_format return f"{path}:{middle}", None, file_format diff --git a/tests/unit/test_imaging_formats.py b/tests/unit/test_imaging_formats.py index d5751aba..3ea61557 100644 --- a/tests/unit/test_imaging_formats.py +++ b/tests/unit/test_imaging_formats.py @@ -23,7 +23,7 @@ import pytest from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import DatasetManagerError -from konfai.utils.utils import SUPPORTED_EXTENSIONS, split_path_spec +from konfai.utils.utils import SUPPORTED_FORMATS, split_path_spec def _image_attributes() -> Attribute: @@ -363,13 +363,13 @@ def test_dicom_write_preserves_unrelated_files(self, tmp_path: Path) -> None: def test_ome_zarr_format_aliases(self, tmp_path: Path, file_format: str) -> None: assert Dataset(tmp_path / file_format, file_format).file_format == "omezarr" - @pytest.mark.parametrize("file_format", ["dicom", "omezarr", "ome-zarr", "ome_zarr", "zarr"]) + @pytest.mark.parametrize("file_format", ["dicom", "omezarr", "ome-zarr", "ome_zarr", "zarr", "itktransform"]) def test_data_manager_path_parser_accepts_imaging_backend(self, file_format: str) -> None: - assert file_format in SUPPORTED_EXTENSIONS + assert file_format in SUPPORTED_FORMATS assert split_path_spec( f"./Dataset:a:{file_format}", allowed_flags={"a", "i"}, - supported_extensions=SUPPORTED_EXTENSIONS, + supported_formats=SUPPORTED_FORMATS, ) == ("./Dataset", "a", file_format) @pytest.mark.parametrize("file_format", ["dicom", "omezarr"]) diff --git a/tests/unit/test_itk_transform_backend.py b/tests/unit/test_itk_transform_backend.py index 216f5781..8530a3c2 100644 --- a/tests/unit/test_itk_transform_backend.py +++ b/tests/unit/test_itk_transform_backend.py @@ -30,6 +30,7 @@ pytest.importorskip("h5py") from konfai.utils.dataset import Attribute, Dataset # noqa: E402 +from konfai.utils.errors import DatasetManagerError # noqa: E402 _ORIGIN, _SPACING = [7.0, -3.0, 10.0], [1.5, 1.5, 2.0] _DIRECTION = [0.0, -1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0] @@ -135,3 +136,36 @@ def test_a_region_read_decodes_only_its_rows_and_matches_the_whole(tmp_path: Pat region, _ = dataset.read_data_slice("Transform", "P000", (slice(0, 3), slice(1, 3), slice(2, 5), slice(0, 4))) np.testing.assert_array_equal(np.asarray(region), np.asarray(whole)[:, 1:3, 2:5, 0:4]) + + +def test_a_transform_dataset_resolves_as_a_run_input(tmp_path: Path) -> None: + """What a run writes, a run can read back. + + ``itktransform`` is a backend token, not a suffix — an entry is ``.h5``, and no path is + ever named ``.itktransform``. Validating a spec against the extensions alone rejected the very + format the write side had just produced. + """ + from konfai.data.data_manager import DataPrediction, Group, GroupTransform + + root = tmp_path / "out" + Dataset(root, "itktransform").write("Transform", "P000", _field(4), _attributes()) + + prediction = DataPrediction( + augmentations=None, + dataset_filenames=[f"{root}:a:itktransform"], + groups_src={ + "Transform": Group(groups_dest={"Transform": GroupTransform(transforms=None, patch_transforms=None)}) + }, + ) + + assert prediction._resolve_dataset_sources() == {"Transform": [(str(root), True)]} + + +def test_without_h5py_the_backend_names_the_extra_to_install(tmp_path: Path, monkeypatch) -> None: + """The parameters are touched region by region through h5py, so it is a requirement, not a + preference: a peak memory that turns on whether an optional import succeeded is one nobody can + size. The ``h5`` backend has always demanded it — this says so out loud.""" + monkeypatch.setattr("konfai.utils.dataset.h5py", None) + + with pytest.raises(DatasetManagerError, match="h5py"): + Dataset(tmp_path / "out", "itktransform").write("Transform", "P000", _field(), _attributes()) From 4337e9cb996f76c692af1bad455966e0f899882b Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 15:29:00 +0200 Subject: [PATCH 36/39] fix(data): a text transform and a stepped region read through the backend get_infos opened every entry as HDF5; a legacy text .tfm is not one, and now takes the whole-decode path. The leading-axis step was dropped by the span read -- rows start..stop came back whole -- and now subsamples the reshaped block; a reversed axis falls back to the whole read. Both pinned by tests, with a foreign text file among them. --- konfai/utils/dataset.py | 22 ++++++++++++++--- tests/unit/test_itk_transform_backend.py | 30 ++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index a0c8378e..e753606e 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -2055,16 +2055,27 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) - region — the peak is the row span, never the field. """ shape, attributes = self.get_infos(group, name) - if len(shape) != 4 or shape[0] != 3 or DISPLACEMENT_FIELD_ATTRIBUTE not in attributes: + leading = slices[1].indices(shape[1]) if len(shape) == 4 else (0, 0, 1) + if ( + len(shape) != 4 + or shape[0] != 3 + or DISPLACEMENT_FIELD_ATTRIBUTE not in attributes + or not h5py.is_hdf5(self._path(name)) + or leading[2] < 0 # a reversed leading axis has no contiguous span to read + ): data, attributes = self.file_to_data(group, name) return data[slices], attributes spatial = shape[1:] - leading = slices[1].indices(spatial[0]) row = 3 * int(np.prod(spatial[1:], dtype=np.int64)) with h5py.File(self._path(name), "r") as file: span = file["TransformGroup/0/TransformParameters"][leading[0] * row : leading[1] * row] block = np.moveaxis(span.reshape(leading[1] - leading[0], *spatial[1:], 3), -1, 0) - return np.asarray(block[(slices[0], slice(None), *slices[2:])], dtype=np.float32), attributes + # The span is the axis WITHOUT its step: rows start..stop were read whole, so the step + # subsamples here, on the reshaped block. + return ( + np.asarray(block[(slices[0], slice(None, None, leading[2]), *slices[2:])], dtype=np.float32), + attributes, + ) def file_to_data_statistics( self, @@ -2133,6 +2144,11 @@ def is_exist(self, group: str, name: str | None = None) -> bool: return os.path.exists(self._path(name if name else group)) def get_infos(self, group: str, name: str) -> tuple[list[int], Attribute]: + # A legacy TEXT transform (`#Insight Transform File V1.0`) is served by the read side + # too; only a real HDF5 file has the parameter datasets this fast path opens. + if not h5py.is_hdf5(self._path(name)): + data, attributes = self.file_to_data(group, name) + return [int(extent) for extent in data.shape], attributes with h5py.File(self._path(name), "r") as file: kind = bytes(file["TransformGroup/0/TransformType"][0]) fixed = np.asarray(file["TransformGroup/0/TransformFixedParameters"][()], dtype=np.float64) diff --git a/tests/unit/test_itk_transform_backend.py b/tests/unit/test_itk_transform_backend.py index 8530a3c2..e9925ff4 100644 --- a/tests/unit/test_itk_transform_backend.py +++ b/tests/unit/test_itk_transform_backend.py @@ -125,6 +125,36 @@ def test_a_foreign_affine_file_reads_back_too(tmp_path: Path) -> None: assert back.TransformPoint(point) == pytest.approx(affine.TransformPoint(point)) +def test_a_foreign_text_transform_serves_headers_without_crashing(tmp_path: Path) -> None: + """A legacy `.tfm` is TEXT (`#Insight Transform File V1.0`), not HDF5: the header fast path + must not open it with h5py.""" + affine = sitk.AffineTransform(3) + affine.SetTranslation((2.0, -1.0, 3.0)) + case = tmp_path / "out" / "P000" + case.mkdir(parents=True) + sitk.WriteTransform(affine, str(case / "Reg.tfm")) + + dataset = Dataset(tmp_path / "out", "itktransform") + shape, _attributes_back = dataset.get_infos("Reg", "P000") + assert len(shape) == 2 # parameter rows, not a field + assert not dataset.bounded_region_reads("Reg", "P000") + back = dataset.read_transform("Reg", "P000") + point = (1.0, 2.0, 3.0) + assert back.TransformPoint(point) == pytest.approx(affine.TransformPoint(point)) + + +def test_a_stepped_region_read_honours_the_leading_axis_step(tmp_path: Path) -> None: + """The row span is read whole and the step subsamples it — same values as slicing the whole.""" + field = _field(5) + dataset = Dataset(tmp_path / "out", "itktransform") + dataset.write("Transform", "P000", field, _attributes()) + + whole, _ = dataset.read_data("Transform", "P000") + region, _ = dataset.read_data_slice("Transform", "P000", (slice(0, 3), slice(0, 4, 2), slice(2, 5), slice(0, 4))) + + np.testing.assert_array_equal(np.asarray(region), np.asarray(whole)[:, 0:4:2, 2:5, 0:4]) + + def test_a_region_read_decodes_only_its_rows_and_matches_the_whole(tmp_path: Path) -> None: """The parameters are HDF5, so a slab reads the span it maps to — same values as the whole.""" field = _field(3) From f7c8c0e109adca45cc085465c8d3c748925a36d9 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 16:13:11 +0200 Subject: [PATCH 37/39] fix(api): spell numpy scalars, and copy a caller's config file np.float64 passed the float check unspelled and ruamel refused it at dump time; the numpy branch now comes first. And a config given as a PATH was rewritten in place by resolution write-back -- the caller's file is not this call's to rewrite, so the write-back lands on a scratch copy, removed at exit. Both pinned. --- konfai/api.py | 27 +++++++++++++++++++++++---- tests/unit/test_api.py | 15 +++++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/konfai/api.py b/konfai/api.py index 12719eff..0b400ee0 100644 --- a/konfai/api.py +++ b/konfai/api.py @@ -37,9 +37,12 @@ rewrite. """ +import atexit import importlib import json import os +import shutil +import tempfile import threading from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager @@ -90,12 +93,13 @@ def _one_workflow_at_a_time(ranks: int) -> Iterator[None]: def _yaml_safe(value: object, where: str) -> object: """``value`` as the config file could hold it — or a refusal that names the argument.""" + # Before the Python scalars: np.float64 IS a float subclass, and ruamel refuses it. + if isinstance(value, np.generic): + return value.item() if value is None or isinstance(value, (bool, int, float, str)): return value if isinstance(value, Path): return str(value) - if isinstance(value, np.generic): - return value.item() if isinstance(value, Mapping): return {str(key): _yaml_safe(entry, f"{where}.{key}") for key, entry in value.items()} if isinstance(value, (list, tuple)): @@ -373,6 +377,21 @@ def evaluate( return EvaluationResult(workspace=workspace, metrics=reports) +def _config_copy(config: "Mapping[str, object] | Path | str") -> "dict[str, object] | Path": + """The caller's config, in a form this call may consume. + + Reading a KonfAI config resolves and REWRITES it -- the record the workspace keeps. A tree is + passed through; a caller's FILE is not this call's to rewrite, so the write-back lands on a + scratch copy instead (removed at exit, like :func:`_materialized_config`'s). + """ + if isinstance(config, Mapping): + return dict(config) + source = Path(config) + scratch = Path(tempfile.mkdtemp(prefix="konfai_config_")) + atexit.register(shutil.rmtree, scratch, ignore_errors=True) + return Path(shutil.copy2(source, scratch / source.name)) + + # ------------------------------------------------------------------------- PREDICTION / TRAINING @@ -397,7 +416,7 @@ def predict( with _one_workflow_at_a_time(len(gpu or []) or cpu): workflow = build_predict( models=[Path(model) for model in models], - prediction_file=config if isinstance(config, (Path, str)) else dict(config), + prediction_file=_config_copy(config), predictions_dir=predictions_dir, ) execute_distributed_object(workflow, gpu=list(gpu or []), cpu=cpu, overwrite=overwrite, quiet=quiet) @@ -430,7 +449,7 @@ def train( workflow = build_train( command=State.RESUME if resume else State.TRAIN, model=model, - config=config if isinstance(config, (Path, str)) else dict(config), + config=_config_copy(config), checkpoints_dir=checkpoints_dir, statistics_dir=statistics_dir, lr=lr, diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py index d9fa5433..b93d4966 100644 --- a/tests/unit/test_api.py +++ b/tests/unit/test_api.py @@ -63,6 +63,21 @@ def test_a_repeated_mapping_stage_is_qualified_by_resolution() -> None: assert list(tree) == ["Clip", "konfai.data.transform:Clip"] +def test_a_numpy_scalar_is_spelled_as_a_plain_scalar() -> None: + """np.float64 IS a float subclass; unspelled, ruamel refuses it at dump time.""" + spelled = api._yaml_safe(np.float64(1.5), "chains.CT.CT.Clip.min_value") + assert type(spelled) is float and spelled == 1.5 + + +def test_a_config_file_is_copied_not_rewritten(tmp_path: Path) -> None: + """Reading a config rewrites it; a caller's file is not this call's to rewrite.""" + source = tmp_path / "Prediction.yml" + source.write_text("Predictor: {}\n", encoding="utf-8") + copy = api._config_copy(source) + assert copy != source + assert Path(copy).read_text(encoding="utf-8") == source.read_text(encoding="utf-8") + + def test_a_subclass_delegating_to_super_keeps_its_own_spelling() -> None: """The recorded spelling is the caller's: a subclass expanding into ``Resample`` arguments inside ``super().__init__`` records its OWN kwargs, so the tree references the subclass with From a77b59bb58f92dddcd34641bf1910357a43a651a Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 16:13:17 +0200 Subject: [PATCH 38/39] fix(data): the backend rewrite lands on .h5, and Std prices its buffers data_to_file resolved an existing .tfm and renamed HDF5 content onto it; ITK selects transform IO from the extension, so that entry was corrupt. The write now always lands on the .h5 name and drops the .tfm it replaces. Std declared working_multiple 0 while holding five buffers beside the region it reads. --- konfai/data/reduction.py | 3 ++ konfai/utils/dataset.py | 37 ++++++++++++++---------- tests/unit/test_itk_transform_backend.py | 15 ++++++++++ 3 files changed, 39 insertions(+), 16 deletions(-) diff --git a/konfai/data/reduction.py b/konfai/data/reduction.py index db944fa9..018fdbfc 100644 --- a/konfai/data/reduction.py +++ b/konfai/data/reduction.py @@ -164,6 +164,9 @@ class Std(Reduction): voxel_local = True incremental = True + # Two persistent accumulators (mean, m2) plus, per accumulate: the float copy of the case, + # ``delta``, and ``value - mean`` again after the mean moved -- five buffers beside the region. + working_multiple = 5.0 def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: self.start() diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index e753606e..8de65016 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -2095,26 +2095,31 @@ def data_to_file( attributes: Attribute | None = None, ) -> None: os.makedirs(self.filename, exist_ok=True) - final = self._path(name) + # Always the `.h5` name: the content is HDF5 and ITK selects its transform IO from the + # extension, so renaming it onto a resolved existing `.tfm` would corrupt that entry. + final = os.path.join(self.filename, f"{name}.h5") staging = f"{self.filename}.{name}.{os.getpid()}.tmp.h5" if isinstance(data, sitk.Transform): sitk.WriteTransform(data, staging) - os.replace(staging, final) - return - if isinstance(data, sitk.Image): - data, attributes = image_to_data(data) - array = np.asarray(data) - if attributes is None or array.ndim != 4 or array.shape[0] != 3: - raise DatasetManagerError( - f"An ':itktransform' entry is a 3-component 3-D displacement field; '{name}' has" - f" shape {list(array.shape)}.", - "Write the field itself (channel-first, with its geometry), or a sitk.Transform.", - ) - spatial = [int(extent) for extent in array.shape[1:]] - file, parameters = _create_itk_transform_file(staging, spatial, attributes) - with file: - parameters[:] = np.moveaxis(array.astype(np.float64), 0, -1).ravel() + else: + if isinstance(data, sitk.Image): + data, attributes = image_to_data(data) + array = np.asarray(data) + if attributes is None or array.ndim != 4 or array.shape[0] != 3: + raise DatasetManagerError( + f"An ':itktransform' entry is a 3-component 3-D displacement field; '{name}' has" + f" shape {list(array.shape)}.", + "Write the field itself (channel-first, with its geometry), or a sitk.Transform.", + ) + spatial = [int(extent) for extent in array.shape[1:]] + file, parameters = _create_itk_transform_file(staging, spatial, attributes) + with file: + parameters[:] = np.moveaxis(array.astype(np.float64), 0, -1).ravel() os.replace(staging, final) + try: # one entry per name: a `.tfm` left under the same stem would double it + os.remove(os.path.join(self.filename, f"{name}.tfm")) + except FileNotFoundError: + pass def open_data_stream( self, diff --git a/tests/unit/test_itk_transform_backend.py b/tests/unit/test_itk_transform_backend.py index e9925ff4..5edeeb5e 100644 --- a/tests/unit/test_itk_transform_backend.py +++ b/tests/unit/test_itk_transform_backend.py @@ -143,6 +143,21 @@ def test_a_foreign_text_transform_serves_headers_without_crashing(tmp_path: Path assert back.TransformPoint(point) == pytest.approx(affine.TransformPoint(point)) +def test_rewriting_a_tfm_entry_lands_on_the_h5_name(tmp_path: Path) -> None: + """The write is HDF5; renamed onto a resolved `.tfm` it would corrupt that entry, since ITK + selects its transform IO from the extension. One entry per name survives the rewrite.""" + affine = sitk.AffineTransform(3) + case = tmp_path / "out" / "P000" + case.mkdir(parents=True) + sitk.WriteTransform(affine, str(case / "Reg.tfm")) + + dataset = Dataset(tmp_path / "out", "itktransform") + dataset.write("Reg", "P000", _field(7), _attributes()) + + assert (case / "Reg.h5").is_file() + assert not (case / "Reg.tfm").exists() + + def test_a_stepped_region_read_honours_the_leading_axis_step(tmp_path: Path) -> None: """The row span is read whole and the step subsamples it — same values as slicing the whole.""" field = _field(5) From 95bc466ff847545c2bf8c279f84e0359a577d600 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 6 Aug 2026 16:13:22 +0200 Subject: [PATCH 39/39] fix(impact-reg): numeric case order past P999, and the Moved name guard sorted() places P1000 before P101 and pairs the wrong moving unit; cases now sort by konfai-apps' own numbering (length, then name). A preset whose output group is named Moved would be deleted by the derivation's own stale-output purge, and is refused by name. --- .../impact_reg_konfai/impact_reg.py | 18 ++++++++++--- .../tests/unit/test_orchestration.py | 27 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/apps/impact_reg/impact_reg_konfai/impact_reg.py b/apps/impact_reg/impact_reg_konfai/impact_reg.py index 5060e2a6..aa0cca54 100644 --- a/apps/impact_reg/impact_reg_konfai/impact_reg.py +++ b/apps/impact_reg/impact_reg_konfai/impact_reg.py @@ -70,6 +70,12 @@ def _as_transform(path: Path) -> "sitk.Transform": return sitk.DisplacementFieldTransform(sitk.Cast(read_displacement_field(path), sitk.sitkVectorFloat64)) +def _case_key(name: str) -> tuple[int, str]: + """konfai-apps numbers cases ``P000``..: zero-padded to three digits, longer past ``P999`` -- + so length-then-lexicographic IS its numeric order, without parsing a preset's own naming.""" + return (len(name), name) + + def _app_id(preset: str) -> str: """Resolve a preset to a KonfAIApp id: a local ``/`` path, or ``:`` on HF.""" if Path(IMPACT_REG_KONFAI_REPO).is_dir(): @@ -443,10 +449,16 @@ def register( " ensemble folds one group, so every member must declare the same one." ) group = groups.pop() + if group == "Moved" and not fields_only: + raise RuntimeError( + "the preset names its output 'Moved', the name this pipeline writes the derived" + " image under; the two would collide in the case directory. Rename the preset's" + " output, or pass fields_only." + ) fields_by_preset = {preset: fields for preset, (_, fields) in fields_by_preset.items()} - cases = sorted(fields_by_preset[presets[0]]) + cases = sorted(fields_by_preset[presets[0]], key=_case_key) for preset, fields in fields_by_preset.items(): - if sorted(fields) != cases: + if sorted(fields, key=_case_key) != cases: raise RuntimeError( f"preset '{preset}' produced cases {sorted(fields)} where '{presets[0]}' produced " f"{cases}; an ensemble can only be averaged case by case." @@ -507,7 +519,7 @@ def _ensemble_mean( from konfai.data.transform import Reduce, Write members = _stage_group( - work / f"ensemble_{case}", "DVF", {preset: dvf for preset, dvf in zip(presets, dvf_paths, strict=True)} + work / f"ensemble_{case}", "DVF", dict(zip(presets, dvf_paths, strict=True)) ) suffixes = "".join(dvf_paths[0].suffixes) _output_path(output / case, group, suffixes) # drop a stale other-form output before writing diff --git a/apps/impact_reg/tests/unit/test_orchestration.py b/apps/impact_reg/tests/unit/test_orchestration.py index d8e59452..64ab9a2c 100644 --- a/apps/impact_reg/tests/unit/test_orchestration.py +++ b/apps/impact_reg/tests/unit/test_orchestration.py @@ -302,6 +302,33 @@ def mixed(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, work, * app.register(["A", "B"], [fixed], [moving], output=tmp_path / "Output") +def test_case_order_is_numeric_past_p999() -> None: + """`sorted` alone puts P1000 before P101 and pairs the wrong moving unit.""" + assert sorted(["P101", "P1000", "P099"], key=reg._case_key) == ["P099", "P101", "P1000"] + + +def test_register_refuses_a_preset_output_named_moved(tmp_path: Path) -> None: + """The derived image is written under 'Moved'; a preset using that name would be deleted by + its own derivation's stale-output purge.""" + moving = tmp_path / "moving.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(moving)) + fixed = tmp_path / "fixed.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((8, 8, 8), dtype=np.float32)), str(fixed)) + + reference = sitk.ReadImage(str(moving)) + app = reg.ImpactRegKonfAIApp() + + def named_moved(preset, fixed_i, moving_i, fixed_masks, moving_masks, n_cases, work, *args, **kwargs): + out = Path(work) / preset / "P000" + out.mkdir(parents=True, exist_ok=True) + _write_dvf(out / "Moved.mha", (2.0, 0.0, 0.0), reference) + return "Moved", {"P000": out / "Moved.mha"} + + app._infer_preset = named_moved # type: ignore[method-assign] + with pytest.raises(RuntimeError, match="Moved"): + app.register(["FireANTs_SyN"], [fixed], [moving], output=tmp_path / "Output") + + def test_find_output_group_discovers_the_one_group(tmp_path: Path) -> None: """konfai-apps writes ``///…``; the group is the one directory holding cases.""" (tmp_path / "reg" / "Transform" / "P000").mkdir(parents=True)