feat(transform): resample onto a reference grid, and warp in one pass

docs/source/config_guide/transform.md

@@ -192,6 +192,122 @@ One chain changes its cardinality at most once. Composing the two — augment a
 cohort, then fold it — is two invocations, the second reading the first one's
 output back.
 
+### `Reduce`: one grid for the cohort, first
+
+Folding cases together only means something if they are the *same* voxels, so
+`Reduce` compares the grid each case's chain **lands on**, and `grid:` says how
+strictly:
+
+| `grid:` | Compares | Use it when |
+| --- | --- | --- |
+| `strict` | extent **and** `Spacing` / `Origin` / `Direction` | the cases really do share a space |
+| `shape_only` | extent alone | you know they share one and the headers disagree |
+| `reference:<case>` | extent; that case's header is what the output carries | one member is the one to believe |
+
+A cohort as acquired rarely passes `strict`: extents differ, and origins can
+differ by more than the volumes are wide, because an acquisition's stage
+coordinates are not an anatomical frame. `ResampleToReference` is what makes
+`strict` true rather than waived — it resamples each case onto the grid of a
+**declared reference**, adopting its extent, spacing, origin and direction:
+
+```yaml
+transforms:
+  ResampleToReference: {entry: case_0, group: CT, fill: 0.0}
+  Reduce: {operator: Median, output: template, grid: strict}
+  Write: {dataset: ./Template:mha}
+```
+
+The reference is a stored image, named by `entry` — and by `group` when the
+store holds more than one. It is looked up by entry, not by the case being
+processed, because one grid serves the whole cohort: in the run's own
+`dataset_filenames`, or in a store of its own.
+
+```yaml
+transforms:
+  ResampleToReference: {entry: case_1, group: CT, dataset: ./Raw:mha}
+  Write: {dataset: ./OnTemplate:mha}
+```
+
+`dataset:` takes the same `path[:format]` spec as everywhere else, so the
+reference can live anywhere — which is the atlas loop: point round N+1 at the
+store round N wrote its template into.
+
+#### Through a displacement field, in one interpolation
+
+Add `field:` and the stage becomes the whole of a registration's apply step —
+`sitk.Resample(image, reference_grid, DisplacementFieldTransform(field))`. For
+each voxel of the **target** grid the field is read at that voxel's world
+position, added, and the source sampled **once** at the displaced point:
+
+```yaml
+transforms:
+  ResampleToReference:
+    entry: case_0
+    group: CT
+    field: ./Fields:mha
+    field_group: DVF
+    max_displacement: 4.0
+  Write: {dataset: ./Registered:mha}
+```
+
+Fields stored *beside* the cases — one entry per case, in the same roots — need
+no path at all: `field_group: DVF` on its own finds them.
+
+Doing this as two stages instead (resample onto the grid, then `Warp`) costs
+**two** interpolations, and the second cannot restore the detail the first
+smoothed away. That is not a small effect: on a high-frequency volume the
+second pass moves voxels by a large fraction of the range, which is exactly why
+an atlas rebuilds its appearance from native volumes rather than from resampled
+ones.
+
+The field lives on **its own grid**, usually coarser than either the source or
+the target, and is read where it is asked — it is defined in world units, so a
+field solved at 120 µm moves a volume stored at 30 µm without being upsampled
+first. Outside its own extent the displacement is zero: the transform is the
+identity where the field says nothing, as SimpleITK has it.
+
+`max_displacement` sizes the source region each target region must read, and is
+**checked against every field region actually read** — a field that exceeds
+what it declared raises rather than sampling zeros, which would show up as a
+dark rim around the moved anatomy and nothing else. It takes `auto`, reading
+the bound KonfAI records on a field it writes. With no bound at all the stage
+declares `WHOLE_VOLUME` and says so in the plan, exactly as `Warp` does.
+
+```{note}
+`Warp` still exists and is not this: it adds a displacement on the case's *own*
+grid, which is the shape update of an atlas build, and it neither changes the
+grid nor needs a reference. Use it when the field was solved on the very grid
+it is applied to.
+```
+
+Naming an image rather than fifteen numbers is deliberate. A grid is an extent
+in array order `(Z, Y, X)` plus an origin, a spacing and a direction in physical
+`(x, y, z)` — transcribing those by hand is the mistake that actually gets made,
+and a transposed grid resamples perfectly well onto the wrong place. A header
+cannot make that mistake.
+
+```{note}
+It streams: a slab of the output reads only the part of the input under it, so a
+case never has to fit in memory. The sampler is `sitk.Resample`'s — linear with
+taps clamped to the buffer, nearest by round-half-up, and `fill` wherever the
+reference grid reaches past the case.
+```
+
+**What it refuses**, rather than write something plausible and wrong:
+
+- a case or a reference carrying no `Origin` / `Spacing` / `Direction` — without
+  geometry there is no physical space to resample in, and a size ratio must not
+  quietly stand in for one;
+- a reference whose `Direction` differs from the case's — the map between them
+  is then a rotation, not a scale and a shift per axis. Reorient first
+  (`Canonical`);
+- a case that does not meet the reference grid **anywhere** — its output would
+  be `fill` from edge to edge, and a median would take that as anatomy.
+
+Partial overlap is legal and ordinary: the rest of the output is `fill`, and the
+plan prints the fraction of the grid each case covers. "Most of this template is
+background" is then something read before the run rather than after it.
+
 ### `Expand`: one case, N copies
 
 `Expand` multiplies, and nothing else. The draws are **ordinary stages of the
@@ -341,32 +457,9 @@ not. `Clip` then `Normalize` therefore takes the whole-volume path, and the plan
 says so. Reorder the chain, or cut it with a `Save`.
 ```
 
-### `ShapeUpdate`: the shape residual of a displacement field
-
-The shape update of an atlas build: `output = -step * (field - t)`, where `t` is
-the field's per-component spatial mean in world units. Resampling a template
-through the result moves it along the cohort's shape residual, at ANTs'
-gradient step.
-
-```yaml
-transforms:
-  ShapeUpdate:
-    step: 0.25
-  Write:
-    dataset: ./Update:omezarr
-```
-
-`t` is **stripped, not applied**, and that is the whole point of the stage. A
-total field maps template coordinates into each specimen's OWN world frame, so
-its spatial mean is dominated by the frame-to-frame offset, not by any pose error
-of the template. Applying it in full translates the template out of its own grid
-and clips the anatomy; stripping it is what keeps the template anchored — the
-same thing ANTs' `AverageAffineTransformNoRigid` is for.
-
-Because the statistic is declared rather than recomputed per region, a field of
-any size runs region by region: the volume is never assembled. Handed the whole
-volume anyway — a chain that fell back for another reason — the stage takes the
-statistic from the tensor in hand, so both paths leave the same state behind.
+Handed the whole volume anyway — a chain that fell back for another reason — a
+stage should take the statistic from the tensor in hand and record it, so both
+paths leave the same state behind.
 
 ## Writing your own transform
 
@@ -434,7 +527,8 @@ refused (cheaper to load the volume), and the plan says so.
 | a bounded neighbourhood | `HALO`, with `halo=(r,)` | nothing |
 | the volume flipped/permuted | `ORIENTATION` | `stream_region_source()` |
 | a translated sub-box | `CROP` | `stream_region_source()` |
-| a resampled grid | `RESCALE` | inherit from `Resample` |
+| the same box, resampled | `RESCALE` | inherit from `Resample` |
+| another grid entirely | `REGRID` | `stream_region_source()` and `stream_region()` |
 | whole-volume Min/Max/Mean/Std | `GLOBAL_STAT`, with `stat_keys` | nothing |
 | the same, per component | `GLOBAL_STAT`, with `MinPerChannel`/`MaxPerChannel`/`MeanPerChannel`/`StdPerChannel` | nothing |
 | genuinely the whole volume | nothing (the default) | nothing |

konfai/data/data_manager.py

@@ -193,6 +193,11 @@ def _check_patch_transform_locality(transform: Transform, group_src: str, group_
             f"'{name}' resamples its input: applied to one patch it rescales that patch about its own"
             " extent and changes the patch grid predictions are reassembled onto."
         ),
+        LocalityKind.REGRID: (
+            f"'{name}' resamples its input onto another grid: applied to one patch it would hand back"
+            " the whole reference extent, which is neither the patch nor the patch grid predictions"
+            " are reassembled onto."
+        ),
         LocalityKind.WHOLE_VOLUME: f"'{name}' needs the whole volume.",
     }
     raise ConfigError(

konfai/data/patching.py

@@ -208,7 +208,13 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute)
 
 
 # The region kinds a composed streamed read (or write) carries between its pointwise stages.
-_REGION_KINDS = (LocalityKind.HALO, LocalityKind.ORIENTATION, LocalityKind.CROP, LocalityKind.RESCALE)
+_REGION_KINDS = (
+    LocalityKind.HALO,
+    LocalityKind.ORIENTATION,
+    LocalityKind.CROP,
+    LocalityKind.RESCALE,
+    LocalityKind.REGRID,
+)
 
 
 # The pull maps are callable dataclasses, not closures, because a plan crosses a process boundary: