From f3b940e6b5d551925b86d5dd6ebdada248297acc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:05:46 +0100 Subject: [PATCH 01/11] Store exact history per batch element Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- docs/get-started/migration.md | 6 +- src/torchio/data/batch.py | 161 +++---------------- src/torchio/data/batch_history.py | 145 +++++++++++++++++ src/torchio/transforms/compose.py | 39 +---- src/torchio/transforms/cornucopia_adapter.py | 1 - src/torchio/transforms/inverse.py | 11 +- src/torchio/transforms/monai_adapter.py | 1 - src/torchio/transforms/transform.py | 91 ++++++----- tests/conftest.py | 15 +- tests/test_anisotropy.py | 6 +- tests/test_batch.py | 6 +- tests/test_bias_field.py | 6 +- tests/test_blur.py | 6 +- tests/test_flip.py | 7 +- tests/test_gamma.py | 9 +- tests/test_ghosting.py | 6 +- tests/test_labels_to_image.py | 18 +-- tests/test_motion.py | 6 +- tests/test_noise.py | 9 +- tests/test_normalize.py | 15 +- tests/test_one_of.py | 11 +- tests/test_per_instance.py | 24 ++- tests/test_spatial.py | 12 +- tests/test_spike.py | 6 +- tests/test_swap.py | 6 +- 25 files changed, 312 insertions(+), 311 deletions(-) create mode 100644 src/torchio/data/batch_history.py diff --git a/docs/get-started/migration.md b/docs/get-started/migration.md index 99c5a9652..c4b61cf48 100644 --- a/docs/get-started/migration.md +++ b/docs/get-started/migration.md @@ -360,7 +360,7 @@ for the capability contract and stochastic-realisation caveats. Prefer vectorized operations on 5D tensors or metadata lists. If logic must call a subject-oriented external API, the current low-level escape hatch is to unbatch, process every subject without changing its schema, -restack, and adopt the prior history: +and restack. Exact per-element histories are preserved automatically: ```python from typing import Any @@ -386,9 +386,7 @@ class StripIdentifier(tio.Transform): for subject in subjects: identifier = subject.metadata["identifier"] subject.metadata["identifier"] = identifier.strip() - rebuilt = tio.SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(batch, subjects) - return rebuilt + return tio.SubjectsBatch.from_subjects(subjects) subject = tio.Subject( diff --git a/src/torchio/data/batch.py b/src/torchio/data/batch.py index 6aea3579d..15035df1c 100644 --- a/src/torchio/data/batch.py +++ b/src/torchio/data/batch.py @@ -3,7 +3,6 @@ from __future__ import annotations import copy as _copy -import dataclasses from collections.abc import Sequence from typing import TYPE_CHECKING from typing import Any @@ -13,13 +12,13 @@ from typing_extensions import Self from .affine import AffineMatrix +from .batch_history import _BatchedHistoryMixin from .batch_schema import _ImageSchema from .batch_schema import _SubjectSchema from .bboxes import BoundingBoxes from .image import Image from .image import LabelMap from .image import ScalarImage -from .invertible import Invertible from .points import Points if TYPE_CHECKING: @@ -29,7 +28,7 @@ _BATCH_META_KEYS = ("_batch_size", "_batched_keys", "_keep") -class ImagesBatch(Invertible): +class ImagesBatch(_BatchedHistoryMixin): """A batch of images with per-sample affines and private prototypes. Wraps a 5D tensor `(B, C, I, J, K)` and a list of `AffineMatrix` @@ -57,6 +56,7 @@ def _initialize( data: Tensor, affines: Sequence[AffineMatrix], prototypes: Sequence[Image], + histories: Sequence[Sequence[Any]] | None = None, ) -> None: """Initialize a validated image batch.""" if data.ndim != 5: @@ -74,7 +74,7 @@ def _initialize( self._data = data self._affines = [affine.clone() for affine in affines] self._prototypes = list(prototypes) - self.applied_transforms: list[Any] = [] + self._initialize_histories(data.shape[0], histories) @classmethod def _from_parts( @@ -82,10 +82,11 @@ def _from_parts( data: Tensor, affines: Sequence[AffineMatrix], prototypes: Sequence[Image], + histories: Sequence[Sequence[Any]] | None = None, ) -> Self: """Build an image batch from validated internal parts.""" batch = cls.__new__(cls) - batch._initialize(data, affines, prototypes) + batch._initialize(data, affines, prototypes, histories) return batch @classmethod @@ -139,7 +140,8 @@ def from_images(cls, images: Sequence[Image]) -> Self: stacked = torch.stack(tensors) affines = [image.affine for image in images] prototypes = [_make_image_prototype(image) for image in images] - return cls._from_parts(stacked, affines, prototypes) + histories = [image.applied_transforms for image in images] + return cls._from_parts(stacked, affines, prototypes, histories) @property def data(self) -> Tensor: @@ -210,10 +212,7 @@ def __getitem__(self, index: int) -> Image: affine=self._affines[index].clone(), ) image._metadata = _copy.deepcopy(prototype.metadata) - image.applied_transforms = [ - *prototype.applied_transforms, - *self.applied_transforms, - ] + image.applied_transforms = list(self.history(index)) return image def __len__(self) -> int: @@ -223,6 +222,10 @@ def unbatch(self) -> list[Image]: """Split the batch into individual images.""" return [self[i] for i in range(self.batch_size)] + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild an image batch from images.""" + return type(self).from_images(items) + @property def has_annotations(self) -> bool: """Whether any image prototype carries annotations.""" @@ -237,7 +240,7 @@ def __repr__(self) -> str: return f"ImagesBatch({cls}, batch_size={b}, shape=({c}, {i}, {j}, {k}))" -class SubjectsBatch(Invertible): +class SubjectsBatch(_BatchedHistoryMixin): """A batch of image columns and per-element object stores. Each image field becomes an `ImagesBatch`. Metadata, points, and @@ -271,32 +274,7 @@ def __init__( self._metadata, ) self._schema: _SubjectSchema | None = None - self.applied_transforms: list[Any] = [] - # When per-element branching occurs (e.g. per-instance OneOf), - # this stores the frozen per-element history prefix. Transforms - # applied afterwards still append to `applied_transforms`, and - # `unbatch()` merges the prefix with the sliced suffix. - self._per_element_history: list[list[Any]] | None = None - - def set_per_element_history(self, histories: list[list[Any]]) -> None: - """Freeze a distinct transform history for each batch element. - - Used when different elements receive different transforms (for - example per-instance [`OneOf`][torchio.OneOf]). Resets the shared - `applied_transforms` so that subsequent transforms accumulate as - a common suffix. - - Args: - histories: One history list per batch element. - """ - if len(histories) != self.batch_size: - msg = ( - f"Expected {self.batch_size} per-element histories," - f" got {len(histories)}" - ) - raise ValueError(msg) - self._per_element_history = [list(history) for history in histories] - self.applied_transforms = [] + self._initialize_histories(self._batch_size) @classmethod def from_subjects(cls, subjects: Sequence[Any]) -> Self: @@ -316,6 +294,9 @@ def from_subjects(cls, subjects: Sequence[Any]) -> Self: metadata=_collect_subject_metadata(subjects, schema), ) batch._schema = schema + batch._set_histories( + [subject.applied_transforms for subject in subjects], + ) return batch @property @@ -448,76 +429,17 @@ def unbatch(self) -> list[Any]: for key, values in self._metadata.items(): kwargs[key] = _copy.deepcopy(values[i]) sub = Subject(**kwargs) - suffix = _slice_history(self.applied_transforms, i) - if self._per_element_history is not None: - sub.applied_transforms = list(self._per_element_history[i]) + suffix - else: - sub.applied_transforms = suffix + sub.applied_transforms = list(self.history(i)) subjects.append(sub) return subjects + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild a subject batch from subjects.""" + return type(self).from_subjects(items) + def __len__(self) -> int: return self.batch_size - def adopt_history(self, source: SubjectsBatch, subjects: list[Any]) -> None: - """Carry transform history from *source* after rebuilding the batch. - - Used by code that unbatches, processes, and re-stacks subjects - (for example the MONAI and Cornucopia adapters). Preserves a - per-element history if *source* had one, otherwise copies the - shared history. - - Args: - source: The batch the subjects were unbatched from. - subjects: The processed subjects, in batch order. - """ - if source._per_element_history is not None: - self.set_per_element_history([s.applied_transforms for s in subjects]) - else: - self.applied_transforms = list(source.applied_transforms) - - def clear_history(self) -> None: - """Remove all applied transform records, including per-element ones.""" - self.applied_transforms = [] - self._per_element_history = None - - def get_inverse_transform(self, **kwargs: Any) -> Any: - """Build a transform that inverts the recorded history. - - Raises: - RuntimeError: If the batch carries per-element histories (from - a per-instance `OneOf`/`SomeOf`), since a single batch - inverse is ambiguous. Call `apply_inverse_transform` - (which inverts each element) or `unbatch()` and invert - each subject. - """ - if self._per_element_history is not None: - msg = ( - "This batch has per-element transform histories from a" - " per-instance OneOf/SomeOf, so a single batch inverse is" - " ambiguous. Call apply_inverse_transform() (which inverts" - " each element) or unbatch() and invert each subject." - ) - raise RuntimeError(msg) - return super().get_inverse_transform(**kwargs) - - def apply_inverse_transform(self, **kwargs: Any) -> SubjectsBatch: - """Apply the inverse of the recorded history. - - When the batch carries per-element histories, each element is - inverted independently and the results are re-stacked. - - Args: - **kwargs: Forwarded to `get_inverse_transform`. - - Returns: - A batch with the transforms undone. - """ - if self._per_element_history is not None: - inverted = [s.apply_inverse_transform(**kwargs) for s in self.unbatch()] - return type(self).from_subjects(inverted) - return super().apply_inverse_transform(**kwargs) - def __repr__(self) -> str: fields = [] for label, store in ( @@ -704,40 +626,3 @@ def _slice_params( else: sliced[key] = value return sliced - - -def _slice_history(history: list[Any], index: int) -> list[Any]: - """Build the per-subject transform history for batch element *index*. - - Batch-shared traces are copied unchanged. Per-instance traces are - sliced to the element's own parameters, and traces whose per-element - keep mask excludes this element are dropped. - - Args: - history: The batch-level list of `AppliedTransform` records. - index: The batch element whose history to build. - - Returns: - The list of `AppliedTransform` records for the element. - """ - sliced: list[Any] = [] - for trace in history: - params = getattr(trace, "params", None) - if not isinstance(params, dict) or "_batched_keys" not in params: - sliced.append(trace) - continue - expected_size = params.get("_batch_size") - if expected_size is not None and not 0 <= index < expected_size: - msg = ( - f"Cannot extract per-instance history for element {index}:" - f" the transform was recorded for a batch of size" - f" {expected_size}" - ) - raise IndexError(msg) - keep = params.get("_keep") - if keep is not None and not keep[index]: - continue - batched_keys = params["_batched_keys"] - new_params = _slice_params(params, index, batched_keys) - sliced.append(dataclasses.replace(trace, params=new_params)) - return sliced diff --git a/src/torchio/data/batch_history.py b/src/torchio/data/batch_history.py new file mode 100644 index 000000000..df12a6fac --- /dev/null +++ b/src/torchio/data/batch_history.py @@ -0,0 +1,145 @@ +"""Exact per-element history support for batch containers.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from typing_extensions import Self + +from .invertible import Invertible + + +class _BatchedHistoryMixin(Invertible): + """Store one exact transform history per batch element.""" + + _histories: list[list[Any]] + + @property + def batch_size(self) -> int: + """Number of batch elements.""" + raise NotImplementedError + + def unbatch(self) -> list[Any]: + """Return individual batch elements.""" + raise NotImplementedError + + def _initialize_histories( + self, + batch_size: int, + histories: Sequence[Sequence[Any]] | None = None, + ) -> None: + """Initialize exact per-element histories.""" + if histories is None: + self._histories = [[] for _ in range(batch_size)] + return + self._set_histories(histories) + + @property + def histories(self) -> tuple[tuple[Any, ...], ...]: + """Immutable view of every element's exact history.""" + return tuple(tuple(history) for history in self._histories) + + def history(self, index: int) -> tuple[Any, ...]: + """Return one element's exact history. + + Args: + index: Batch element index. + + Returns: + Immutable transform-history view for the element. + """ + return tuple(self._histories[index]) + + @property + def has_divergent_history(self) -> bool: + """Whether element histories differ.""" + first = self._histories[0] + return any( + not _histories_equal(first, history) for history in self._histories[1:] + ) + + @property + def applied_transforms(self) -> list[Any]: + """Uniform batch history for compatibility. + + Raises: + RuntimeError: If element histories differ. + """ + if self.has_divergent_history: + msg = ( + "This batch has divergent element histories. Use `histories`" + " or `history(index)` instead of `applied_transforms`." + ) + raise RuntimeError(msg) + return list(self._histories[0]) + + @applied_transforms.setter + def applied_transforms(self, history: Sequence[Any]) -> None: + copied = list(history) + self._histories = [list(copied) for _ in range(self.batch_size)] + + def _set_histories(self, histories: Sequence[Sequence[Any]]) -> None: + """Replace every element history.""" + if len(histories) != self.batch_size: + msg = f"Expected {self.batch_size} histories, got {len(histories)}" + raise ValueError(msg) + self._histories = [list(history) for history in histories] + + def _append_history(self, traces: Sequence[Any | None]) -> None: + """Append one optional trace per batch element.""" + if len(traces) != self.batch_size: + msg = f"Expected {self.batch_size} traces, got {len(traces)}" + raise ValueError(msg) + for history, trace in zip(self._histories, traces, strict=True): + if trace is not None: + history.append(trace) + + def clear_history(self) -> None: + """Remove every element history.""" + self._histories = [[] for _ in range(self.batch_size)] + + def get_inverse_transform(self, **kwargs: Any) -> Any: + """Build a vectorized inverse for a uniform batch history. + + Args: + **kwargs: Forwarded to `Invertible.get_inverse_transform`. + + Raises: + RuntimeError: If element histories differ. + """ + if self.has_divergent_history: + msg = ( + "This batch has divergent element histories, so one vectorized" + " inverse is ambiguous. Use `apply_inverse_transform()`." + ) + raise RuntimeError(msg) + return super().get_inverse_transform(**kwargs) + + def apply_inverse_transform(self, **kwargs: Any) -> Self: + """Apply vectorized or per-element inverse transforms. + + Args: + **kwargs: Forwarded to `get_inverse_transform`. + + Returns: + A batch with transforms undone. + """ + if not self.has_divergent_history: + return super().apply_inverse_transform(**kwargs) + inverted = [item.apply_inverse_transform(**kwargs) for item in self.unbatch()] + result = self._batch_items(inverted) + result.clear_history() + return result + + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild the concrete batch type from unbatched items.""" + raise NotImplementedError + + +def _histories_equal(first: Sequence[Any], second: Sequence[Any]) -> bool: + """Compare histories, treating ambiguous value equality as unequal.""" + try: + return bool(first == second) + except (RuntimeError, TypeError, ValueError): + return False diff --git a/src/torchio/transforms/compose.py b/src/torchio/transforms/compose.py index 68d9a97be..a7db004de 100644 --- a/src/torchio/transforms/compose.py +++ b/src/torchio/transforms/compose.py @@ -283,9 +283,7 @@ def to_hydra(self) -> dict[str, Any]: def _apply_to_element(subject: Any, apply_fn: Any) -> Any: """Apply a transform (or callable) to one element, preserving history. - Wrapping a `Subject` directly into a batch discards its existing - history, so the element is wrapped into a one-element batch seeded - with its prior history; the transform then appends to that history. + The one-element batch inherits the subject's exact history. Args: subject: The single subject to transform (carrying its history). @@ -298,7 +296,6 @@ def _apply_to_element(subject: Any, apply_fn: Any) -> Any: from ..data.batch import SubjectsBatch element_batch = SubjectsBatch.from_subjects([subject]) - element_batch.applied_transforms = list(subject.applied_transforms) element_batch = apply_fn(element_batch) return element_batch.unbatch()[0] @@ -317,10 +314,9 @@ def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any: """ from ..data.batch import SubjectsBatch - _check_consistent_schema(subjects, transform_name) try: batch = SubjectsBatch.from_subjects(subjects) - except (RuntimeError, KeyError) as error: + except ValueError as error: msg = ( f"Per-instance {transform_name} produced batch elements with" " different shapes or schemas, which cannot be re-stacked. Use" @@ -328,35 +324,4 @@ def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any: f" {transform_name}, or pass per_instance=False." ) raise RuntimeError(msg) from error - batch.set_per_element_history([s.applied_transforms for s in subjects]) return batch - - -def _check_consistent_schema(subjects: list[Any], transform_name: str) -> None: - """Ensure all subjects share the same image names and classes. - - Per-element branching may apply different transforms to different - elements; if those change the set of images (or their type), the - elements can no longer be re-stacked into one batch. This raises a - clear error instead of silently dropping data. - - Args: - subjects: The subjects about to be re-stacked. - transform_name: Name of the branching transform for the message. - - Raises: - RuntimeError: If image names or classes differ across subjects. - """ - if not subjects: - return - reference = {name: type(image) for name, image in subjects[0].images.items()} - for subject in subjects[1:]: - current = {name: type(image) for name, image in subject.images.items()} - if current != reference: - msg = ( - f"Per-instance {transform_name} produced batch elements with" - " different image names or types, which cannot be re-stacked." - " Use only schema-preserving transforms with per-instance" - f" {transform_name}, or pass per_instance=False." - ) - raise RuntimeError(msg) diff --git a/src/torchio/transforms/cornucopia_adapter.py b/src/torchio/transforms/cornucopia_adapter.py index f8b0588d7..1edef558f 100644 --- a/src/torchio/transforms/cornucopia_adapter.py +++ b/src/torchio/transforms/cornucopia_adapter.py @@ -76,7 +76,6 @@ def forward(self, data: Any) -> Any: from ..data.batch import SubjectsBatch result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) return unwrap(result) def apply_transform( diff --git a/src/torchio/transforms/inverse.py b/src/torchio/transforms/inverse.py index 56bff7ba2..1efd4d960 100644 --- a/src/torchio/transforms/inverse.py +++ b/src/torchio/transforms/inverse.py @@ -80,15 +80,16 @@ def apply_inverse_transform( Returns: Data with transforms undone, same type as input. """ - if not hasattr(data, "applied_transforms"): - return data - # Batches with per-element histories (from per-instance OneOf/SomeOf) - # know how to invert each element; delegate to their own method. - if getattr(data, "_per_element_history", None) is not None: + from ..data.batch import ImagesBatch + from ..data.batch import SubjectsBatch + + if isinstance(data, (ImagesBatch, SubjectsBatch)): return data.apply_inverse_transform( warn=warn, ignore_intensity=ignore_intensity, ) + if not hasattr(data, "applied_transforms"): + return data inverse = get_inverse_transform( data.applied_transforms, warn=warn, diff --git a/src/torchio/transforms/monai_adapter.py b/src/torchio/transforms/monai_adapter.py index 64cb94b40..21491c775 100644 --- a/src/torchio/transforms/monai_adapter.py +++ b/src/torchio/transforms/monai_adapter.py @@ -90,7 +90,6 @@ def forward(self, data): from ..data.batch import SubjectsBatch result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) return unwrap(result) def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 2fbb4882f..76f097e96 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -21,6 +21,7 @@ from ..data.batch import ImagesBatch from ..data.batch import SubjectsBatch +from ..data.batch import _slice_params from ..data.bboxes import BoundingBoxes from ..data.image import Image from ..data.image import ScalarImage @@ -49,21 +50,6 @@ class AppliedTransform: _TRANSFORM_REGISTRY: dict[str, type[Transform]] = {} -def _all_elements_gated_out(params: dict[str, Any]) -> bool: - """Whether per-element gating masked out every batch element. - - Args: - params: The parameter dict produced by `make_params`, possibly - carrying a `_keep` mask added by `_tag_batched`. - - Returns: - `True` only when a `_keep` mask is present and none of its - elements are kept, i.e. the transform was an exact no-op. - """ - keep = params.get("_keep") - return keep is not None and not any(keep) - - def _copy_optional_list(value: list[str] | None) -> list[str] | None: return None if value is None else list(value) @@ -266,31 +252,58 @@ def forward(self, data: Any) -> Any: if not _all_elements_gated_out(params): self._check_spatial_annotations(batch) batch = self.apply_transform(batch, params) - # Record history on the batch, unless every element was gated out by - # per-element probability: that is an exact no-op, and recording it - # would let history replay (e.g. an invertible spatial transform) - # trigger an unnecessary identity resample. - if not _all_elements_gated_out(params): - trace = AppliedTransform( - name=type(self).__name__, - params=params, - include=_copy_optional_list(self.include), - exclude=_copy_optional_list(self.exclude), - ) - if not hasattr(batch, "applied_transforms"): - batch.applied_transforms = [] - batch.applied_transforms.append(trace) + batch._append_history(self._build_history_traces(params, batch.batch_size)) result = unwrap(batch) # Propagate history to outputs that can carry it - if ( - hasattr(batch, "applied_transforms") - and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray)) - and not isinstance(result, dict) - ): + if not isinstance( + result, + (ImagesBatch, SubjectsBatch, Tensor, np.ndarray), + ) and not isinstance(result, dict): with contextlib.suppress(AttributeError): result.applied_transforms = list(batch.applied_transforms) return result + def _build_history_traces( + self, + params: dict[str, Any], + batch_size: int, + ) -> list[AppliedTransform | None]: + """Build one clean optional history trace per element.""" + batched_keys = params.get("_batched_keys") + if batched_keys is None: + return [ + self._make_applied_transform(_copy.deepcopy(params)) + for _ in range(batch_size) + ] + expected_size = params.get("_batch_size") + if expected_size != batch_size: + msg = ( + f"Parameter batch size {expected_size} does not match" + f" input batch size {batch_size}" + ) + raise ValueError(msg) + keep = params.get("_keep") + traces: list[AppliedTransform | None] = [] + for index in range(batch_size): + if keep is not None and not keep[index]: + traces.append(None) + continue + element_params = _slice_params(params, index, batched_keys) + traces.append(self._make_applied_transform(element_params)) + return traces + + def _make_applied_transform( + self, + params: dict[str, Any], + ) -> AppliedTransform: + """Build one history trace.""" + return AppliedTransform( + name=type(self).__name__, + params=params, + include=_copy_optional_list(self.include), + exclude=_copy_optional_list(self.exclude), + ) + def _check_spatial_annotations(self, data: Any) -> None: """Reject spatial transforms that would leave stale annotations.""" if isinstance(self, SpatialTransform) and _data_has_annotations(data): @@ -545,6 +558,7 @@ def _wrap( return data, _unwrap_subjects_batch case ImagesBatch(): sb = SubjectsBatch({"tio_default_image": data}) + sb._set_histories(data.histories) return sb, _unwrap_images_batch case Subject(): sb = SubjectsBatch.from_subjects([data]) @@ -560,6 +574,7 @@ def _wrap_single_image(img: Image, unwrap_fn: Any) -> tuple[Any, Any]: from ..data.batch import SubjectsBatch sub = Subject(tio_default_image=img) + sub.applied_transforms = list(img.applied_transforms) sb = SubjectsBatch.from_subjects([sub]) return sb, unwrap_fn @@ -655,7 +670,9 @@ def _unwrap_subjects_batch(batch: SubjectsBatch) -> SubjectsBatch: def _unwrap_images_batch(batch: SubjectsBatch) -> ImagesBatch: - return batch.images["tio_default_image"] + image_batch = batch.images["tio_default_image"] + image_batch._set_histories(batch.histories) + return image_batch def _unwrap_subject(batch: SubjectsBatch) -> Subject: @@ -664,7 +681,9 @@ def _unwrap_subject(batch: SubjectsBatch) -> Subject: def _unwrap_image(batch: SubjectsBatch) -> Image: sub = batch.unbatch()[0] - return sub.tio_default_image + image = sub.tio_default_image + image.applied_transforms = list(sub.applied_transforms) + return image def _unwrap_tensor(batch: SubjectsBatch) -> Tensor: diff --git a/tests/conftest.py b/tests/conftest.py index 95cd5456d..9139ea52f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,6 @@ import torch from torchio.data.batch import SubjectsBatch -from torchio.data.batch import _slice_params @pytest.fixture @@ -38,10 +37,6 @@ def _assert( ) -> None: original = copy.deepcopy(batch) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params, "per-instance path was not active" - batched_keys = params["_batched_keys"] - keep = params.get("_keep") image_names = list(transform._get_images(result).keys()) result_images = transform._get_images(result) original_subjects = original.unbatch() @@ -51,10 +46,14 @@ def _assert( name: image.data.clone() for name, image in transform._get_images(single).items() } - element_params = _slice_params(params, index, batched_keys) - single = transform.apply_transform(single, element_params) + history = result.history(index) + gated_out = not history + if history: + single = transform.apply_transform( + single, + history[-1].params, + ) single_images = transform._get_images(single) - gated_out = keep is not None and not keep[index] for name in image_names: result_row = result_images[name].data[index : index + 1] torch.testing.assert_close( diff --git a/tests/test_anisotropy.py b/tests/test_anisotropy.py index 5692b28cd..5b4f17bb5 100644 --- a/tests/test_anisotropy.py +++ b/tests/test_anisotropy.py @@ -63,9 +63,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Anisotropy(downsampling=(2.0, 5.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["factor"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["factor"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_batch.py b/tests/test_batch.py index e9331cdc0..c7a6f04fc 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -566,14 +566,11 @@ def _batch(self, batch_size: int = 4) -> SubjectsBatch: return SubjectsBatch.from_subjects(subjects) def test_adopt_history_preserves_per_element(self) -> None: - # Simulate the adapter pattern: a per-element batch is unbatched, - # processed, and re-stacked; history must survive. torch.manual_seed(0) batch = self._batch() branched = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) subjects = branched.unbatch() rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(branched, subjects) for original, restored in zip( branched.unbatch(), rebuilt.unbatch(), @@ -589,7 +586,6 @@ def test_adopt_history_shared_case(self) -> None: transformed = tio.Gamma(log_gamma=0.3, per_instance=False)(batch) subjects = transformed.unbatch() rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(transformed, subjects) - assert rebuilt._per_element_history is None + assert not rebuilt.has_divergent_history for subject in rebuilt.unbatch(): assert [t.name for t in subject.applied_transforms] == ["Gamma"] diff --git a/tests/test_bias_field.py b/tests/test_bias_field.py index 9167fd7f1..f09346479 100644 --- a/tests/test_bias_field.py +++ b/tests/test_bias_field.py @@ -94,9 +94,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.BiasField(std=(0.3, 0.6))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["std"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_shares_std(self) -> None: diff --git a/tests/test_blur.py b/tests/test_blur.py index 101b84fbb..32374f1b1 100644 --- a/tests/test_blur.py +++ b/tests/test_blur.py @@ -50,9 +50,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Blur(std=(1.0, 4.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({tuple(item["std"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) diff --git a/tests/test_flip.py b/tests/test_flip.py index 60b766cdf..26f7736bc 100644 --- a/tests/test_flip.py +++ b/tests/test_flip.py @@ -263,10 +263,9 @@ def test_per_instance_axes_differ_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["axes"]) == batch.batch_size - distinct = {tuple(a) for a in params["axes"]} + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + distinct = {tuple(item["axes"]) for item in params} assert len(distinct) > 1 def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_gamma.py b/tests/test_gamma.py index 8a82abb21..fb062139e 100644 --- a/tests/test_gamma.py +++ b/tests/test_gamma.py @@ -92,10 +92,9 @@ def test_per_instance_default_differs_across_batch(self) -> None: batch = self._batch() transform = tio.Gamma(log_gamma=(0.2, 0.8)) result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(params["log_gamma"]) == batch.batch_size - assert len(set(params["log_gamma"])) > 1 + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] + assert all(isinstance(value, float) for value in log_gammas) + assert len(set(log_gammas)) > 1 def test_per_instance_false_is_shared(self) -> None: torch.manual_seed(0) @@ -117,7 +116,7 @@ def test_per_instance_values_applied(self) -> None: original = batch.t1.data.clone() transform = tio.Gamma(log_gamma=(0.2, 0.8)) result = transform(batch) - log_gammas = result.applied_transforms[-1].params["log_gamma"] + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] for i, log_gamma in enumerate(log_gammas): gamma = torch.tensor(log_gamma).exp() expected = original[i].sign() * original[i].abs().pow(gamma) diff --git a/tests/test_ghosting.py b/tests/test_ghosting.py index 6c97ef233..a0d9dc2b2 100644 --- a/tests/test_ghosting.py +++ b/tests/test_ghosting.py @@ -61,9 +61,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Ghosting(intensity=(0.5, 1.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["intensity"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_labels_to_image.py b/tests/test_labels_to_image.py index 70a0942e4..bad8992a1 100644 --- a/tests/test_labels_to_image.py +++ b/tests/test_labels_to_image.py @@ -72,10 +72,9 @@ def test_per_instance_means_differ_across_batch(self) -> None: batch = self._batch() transform = tio.LabelsToImage(label_key="seg", default_mean=(0.2, 0.9)) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["means"]) == batch.batch_size - means_for_label_1 = [m[1] for m in params["means"]] + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + means_for_label_1 = [item["means"][1] for item in params] assert len(set(means_for_label_1)) > 1 assert result.image_from_labels.data.shape[0] == batch.batch_size @@ -117,18 +116,19 @@ def test_each_element_uses_its_own_label_stats(self) -> None: ) torch.manual_seed(1) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params image = result.img.data for index in range(batch.batch_size): + params = result.history(index)[-1].params region_one = image[index, 0, : size // 2] region_two = image[index, 0, size // 2 :] assert region_one.mean().item() == pytest.approx( - params["means"][index][1], abs=0.5 + params["means"][1], abs=0.5 ) assert region_two.mean().item() == pytest.approx( - params["means"][index][2], abs=0.5 + params["means"][2], abs=0.5 ) # Independent per-element sampling: means vary across the batch. - label_one_means = {round(params["means"][i][1], 3) for i in range(3)} + label_one_means = { + round(result.history(i)[-1].params["means"][1], 3) for i in range(3) + } assert len(label_one_means) > 1 diff --git a/tests/test_motion.py b/tests/test_motion.py index b3c5d1304..5979b610f 100644 --- a/tests/test_motion.py +++ b/tests/test_motion.py @@ -61,9 +61,9 @@ def test_per_instance_differs_across_batch(self) -> None: result = tio.Motion(degrees=(5, 15), translation=(5, 15), num_transforms=2)( batch ) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["transforms"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["transforms"]) for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_noise.py b/tests/test_noise.py index e5c912603..767b0e6fb 100644 --- a/tests/test_noise.py +++ b/tests/test_noise.py @@ -244,10 +244,9 @@ def test_per_instance_default_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Noise(std=(0.5, 1.5))(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["std"], list) - assert len(params["std"]) == batch.batch_size - assert len(set(params["std"])) > 1 + stds = [history[-1].params["std"] for history in result.histories] + assert all(isinstance(value, float) for value in stds) + assert len(set(stds)) > 1 def test_per_instance_false_is_shared(self) -> None: torch.manual_seed(0) @@ -265,7 +264,7 @@ def test_per_instance_mean_applied_per_element(self) -> None: torch.manual_seed(0) batch = self._batch(batch_size=5) result = tio.Noise(mean=(5.0, 20.0), std=0.0)(batch) - means = result.applied_transforms[-1].params["mean"] + means = [history[-1].params["mean"] for history in result.histories] for i, mean in enumerate(means): torch.testing.assert_close( result.t1.data[i].mean(), diff --git a/tests/test_normalize.py b/tests/test_normalize.py index 788e605c0..f4a05b219 100644 --- a/tests/test_normalize.py +++ b/tests/test_normalize.py @@ -278,15 +278,14 @@ def test_per_instance_out_range_differs(self) -> None: batch = self._batch() transform = tio.RescaleIntensity(out_min=(-1.0, 0.0), out_max=(0.5, 1.0)) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["out_min"]) == batch.batch_size - assert len(set(params["out_min"])) > 1 + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["out_min"] for item in params}) > 1 # Each element rescaled to its own output range. for i in range(batch.batch_size): data = result.t1.data[i] - assert data.min() >= params["out_min"][i] - 1e-4 - assert data.max() <= params["out_max"][i] + 1e-4 + assert data.min() >= params[i]["out_min"] - 1e-4 + assert data.max() <= params[i]["out_max"] + 1e-4 def test_per_instance_false_shares_params(self) -> None: torch.manual_seed(0) @@ -314,7 +313,9 @@ def test_per_instance_inverse_zero_range_no_nan(self) -> None: batch = tio.SubjectsBatch.from_subjects(subjects) transform = tio.RescaleIntensity(out_min=0.0, out_max=0.0) result = transform(batch) - assert "_batched_keys" in result.applied_transforms[-1].params + assert all( + "_batched_keys" not in history[-1].params for history in result.histories + ) restored = result.apply_inverse_transform() assert not torch.isnan(restored.t1.data).any() # Identical inputs: the batch-shared input range covers every diff --git a/tests/test_one_of.py b/tests/test_one_of.py index 1d03b5fee..9887e26fa 100644 --- a/tests/test_one_of.py +++ b/tests/test_one_of.py @@ -127,9 +127,10 @@ def test_functional_inverse_restores_per_element(self) -> None: def test_get_inverse_transform_raises_for_per_element(self) -> None: torch.manual_seed(0) - batch = self._batch(batch_size=4) - result = tio.OneOf([tio.Flip(axes=(0,))])(batch) - with pytest.raises(RuntimeError, match="per-element"): + batch = self._batch(batch_size=16) + result = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) + assert result.has_divergent_history + with pytest.raises(RuntimeError, match="divergent"): result.get_inverse_transform() def test_clear_history_clears_per_element(self) -> None: @@ -137,7 +138,7 @@ def test_clear_history_clears_per_element(self) -> None: batch = self._batch(batch_size=4) result = tio.OneOf([tio.Flip(axes=(0,))])(batch) result.clear_history() - assert result._per_element_history is None + assert all(not history for history in result.histories) for subject in result.unbatch(): assert subject.applied_transforms == [] @@ -149,7 +150,7 @@ def test_p_zero_is_noop_preserving_history(self) -> None: flipped = tio.Flip(axes=(0,))(batch) result = tio.OneOf([tio.Flip(axes=(1,))], p=0.0)(flipped) torch.testing.assert_close(result.t1.data, flipped.t1.data) - assert result._per_element_history is None + assert not result.has_divergent_history # The shared Flip history is intact and still invertible as a batch. restored = result.apply_inverse_transform() torch.testing.assert_close(restored.t1.data, original) diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index d5e0bdef6..cb0752bc3 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -82,9 +82,9 @@ def test_compose_child_is_per_instance(self) -> None: batch = _identical_batch() pipeline = tio.Compose([tio.Gamma(log_gamma=(0.2, 0.8))]) result = pipeline(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(set(params["log_gamma"])) > 1 + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] + assert all(isinstance(value, float) for value in log_gammas) + assert len(set(log_gammas)) > 1 def test_compose_respects_per_instance_false(self) -> None: torch.manual_seed(0) @@ -100,7 +100,9 @@ def test_unbatch_slices_history(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Gamma(log_gamma=(0.2, 0.8))(batch) - batch_log_gammas = result.applied_transforms[-1].params["log_gamma"] + batch_log_gammas = [ + history[-1].params["log_gamma"] for history in result.histories + ] for i, subject in enumerate(result.unbatch()): trace = subject.applied_transforms[-1] assert trace.params["log_gamma"] == batch_log_gammas[i] @@ -112,23 +114,17 @@ def test_mismatched_batch_size_raises(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) transform = tio.Affine(degrees=(20.0, 80.0), default_pad_value=0.0) - result = transform(batch) - params = result.applied_transforms[-1].params + params = transform.make_params(batch) smaller = _identical_batch(batch_size=2) with pytest.raises(RuntimeError, match="Per-instance spatial parameters"): transform.apply_transform(smaller, params) - def test_history_slice_out_of_range_raises(self) -> None: - # Slicing per-instance history for an element beyond the recorded - # batch size must fail with a clear error rather than an opaque one. - from torchio.data.batch import _slice_history - + def test_history_index_out_of_range_raises(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Noise(std=(0.1, 0.5))(batch) - history = result.applied_transforms - with pytest.raises(IndexError, match="batch of size 4"): - _slice_history(history, 4) + with pytest.raises(IndexError): + result.history(4) class TestPerInstanceDtypePreservation: diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 4c98800cc..eed80fa3b 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -410,9 +410,9 @@ def test_per_instance_rotations_differ(self) -> None: batch = self._identical_batch() transform = AffineTransform(degrees=(20.0, 80.0), default_pad_value=0.0) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["affine_matrix"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["affine_matrix"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) assert not torch.allclose(data[1], data[2]) @@ -502,9 +502,9 @@ def test_per_instance_elastic_differs_across_batch(self) -> None: max_displacement=(1.0, 3.0), ) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["control_points"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["control_points"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) assert not torch.allclose(data[1], data[2]) diff --git a/tests/test_spike.py b/tests/test_spike.py index 3d1fc1729..aaf31b346 100644 --- a/tests/test_spike.py +++ b/tests/test_spike.py @@ -56,9 +56,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Spike(intensity=(1.0, 3.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["intensity"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_swap.py b/tests/test_swap.py index 45aba36d5..eb09c4905 100644 --- a/tests/test_swap.py +++ b/tests/test_swap.py @@ -63,9 +63,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Swap(patch_size=4, num_iterations=20)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["locations"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["locations"]) for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: From 9a5353f401736a487aaa743cac24d872b448689c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:33:34 +0100 Subject: [PATCH 02/11] Handle prior history in vectorization checks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- tests/conftest.py | 6 ++++-- tests/test_vectorization.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 9139ea52f..32177b9bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -46,9 +46,11 @@ def _assert( name: image.data.clone() for name, image in transform._get_images(single).items() } + prior_history_size = len(original.history(index)) history = result.history(index) - gated_out = not history - if history: + gated_out = len(history) == prior_history_size + if not gated_out: + assert len(history) == prior_history_size + 1 single = transform.apply_transform( single, history[-1].params, diff --git a/tests/test_vectorization.py b/tests/test_vectorization.py index ea7364d82..860c7f892 100644 --- a/tests/test_vectorization.py +++ b/tests/test_vectorization.py @@ -79,3 +79,13 @@ def test_anisotropy_tie_rounding_matches_scalar(assert_vectorized) -> None: [tio.Subject(t1=tio.ScalarImage(data.clone() + index)) for index in range(4)] ) assert_vectorized(tio.Anisotropy(downsampling=2.0), batch) + + +def test_gating_with_prior_history(assert_vectorized) -> None: + torch.manual_seed(0) + batch = tio.Gamma(log_gamma=0.2, per_instance=False)(_batch(batch_size=6)) + + assert_vectorized( + tio.Flip(axes=(0, 1, 2), flip_probability=1.0, p=0.5), + batch, + ) From 7405387b85a05e65d3d424a3262990df4a9f1da4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:40:12 +0100 Subject: [PATCH 03/11] Clarify batch history bounds errors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/data/batch_history.py | 6 ++++++ tests/test_per_instance.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/torchio/data/batch_history.py b/src/torchio/data/batch_history.py index df12a6fac..7524d30d3 100644 --- a/src/torchio/data/batch_history.py +++ b/src/torchio/data/batch_history.py @@ -49,6 +49,12 @@ def history(self, index: int) -> tuple[Any, ...]: Returns: Immutable transform-history view for the element. """ + if not 0 <= index < self.batch_size: + msg = ( + f"Cannot get history for element {index}:" + f" batch size is {self.batch_size}" + ) + raise IndexError(msg) return tuple(self._histories[index]) @property diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index cb0752bc3..25f68c4c4 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -123,7 +123,7 @@ def test_history_index_out_of_range_raises(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Noise(std=(0.1, 0.5))(batch) - with pytest.raises(IndexError): + with pytest.raises(IndexError, match=r"element 4.*batch size is 4"): result.history(4) From a81913998ab3866a78ee191aae0ba49ca0589183 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:53:31 +0100 Subject: [PATCH 04/11] Reuse uniform traces and preserve no-op gating Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/transform.py | 11 +++++------ tests/test_per_instance.py | 8 ++++++++ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 76f097e96..175209a3b 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -249,10 +249,11 @@ def forward(self, data: Any) -> Any: if not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p: return unwrap(batch) params = self.make_params(batch) - if not _all_elements_gated_out(params): + traces = self._build_history_traces(params, batch.batch_size) + if any(trace is not None for trace in traces): self._check_spatial_annotations(batch) batch = self.apply_transform(batch, params) - batch._append_history(self._build_history_traces(params, batch.batch_size)) + batch._append_history(traces) result = unwrap(batch) # Propagate history to outputs that can carry it if not isinstance( @@ -271,10 +272,8 @@ def _build_history_traces( """Build one clean optional history trace per element.""" batched_keys = params.get("_batched_keys") if batched_keys is None: - return [ - self._make_applied_transform(_copy.deepcopy(params)) - for _ in range(batch_size) - ] + trace = self._make_applied_transform(_copy.deepcopy(params)) + return [trace] * batch_size expected_size = params.get("_batch_size") if expected_size != batch_size: msg = ( diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index 25f68c4c4..bbf2f5e31 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -108,6 +108,14 @@ def test_unbatch_slices_history(self) -> None: assert trace.params["log_gamma"] == batch_log_gammas[i] assert "_batched_keys" not in trace.params + def test_uniform_history_reuses_one_trace(self) -> None: + batch = _identical_batch(batch_size=4) + + result = tio.Gamma(log_gamma=0.2, per_instance=False)(batch) + + trace = result.history(0)[-1] + assert all(history[-1] is trace for history in result.histories) + class TestSpatialBatchSizeValidation: def test_mismatched_batch_size_raises(self) -> None: From ec59cf7de30bfae8d82a25880f768a3cabf16ce6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:06:32 +0100 Subject: [PATCH 05/11] Make uniform batch history views immutable Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/data/batch_history.py | 6 +++--- src/torchio/data/image.py | 2 ++ src/torchio/data/invertible.py | 3 ++- src/torchio/data/subject.py | 2 ++ src/torchio/transforms/inverse.py | 3 ++- tests/test_batch.py | 7 +++++++ tests/test_per_instance.py | 2 +- 7 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/torchio/data/batch_history.py b/src/torchio/data/batch_history.py index 7524d30d3..0bbb3cc7e 100644 --- a/src/torchio/data/batch_history.py +++ b/src/torchio/data/batch_history.py @@ -66,8 +66,8 @@ def has_divergent_history(self) -> bool: ) @property - def applied_transforms(self) -> list[Any]: - """Uniform batch history for compatibility. + def applied_transforms(self) -> tuple[Any, ...]: + """Immutable uniform batch history for compatibility. Raises: RuntimeError: If element histories differ. @@ -78,7 +78,7 @@ def applied_transforms(self) -> list[Any]: " or `history(index)` instead of `applied_transforms`." ) raise RuntimeError(msg) - return list(self._histories[0]) + return tuple(self._histories[0]) @applied_transforms.setter def applied_transforms(self, history: Sequence[Any]) -> None: diff --git a/src/torchio/data/image.py b/src/torchio/data/image.py index 000ee4366..eb6b5895f 100644 --- a/src/torchio/data/image.py +++ b/src/torchio/data/image.py @@ -157,6 +157,8 @@ class Image(Invertible): >>> image = tio.ScalarImage(nifti_image) # from nibabel (lazy) """ + applied_transforms: list[Any] + #: Source types accepted by the constructor. ImageInput = ( ImageSource # str | Path | IOBase | OpenFile diff --git a/src/torchio/data/invertible.py b/src/torchio/data/invertible.py index 65a890f6b..34702ae98 100644 --- a/src/torchio/data/invertible.py +++ b/src/torchio/data/invertible.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any from typing_extensions import Self @@ -17,7 +18,7 @@ class Invertible: `self.applied_transforms = []` in their constructor. """ - applied_transforms: list[Any] + applied_transforms: Sequence[Any] def get_inverse_transform( self, diff --git a/src/torchio/data/subject.py b/src/torchio/data/subject.py index fe0d5e543..5f5e323db 100644 --- a/src/torchio/data/subject.py +++ b/src/torchio/data/subject.py @@ -56,6 +56,8 @@ class Subject(Invertible): >>> subject.age # metadata access (returns 45) """ + applied_transforms: list[Any] + def __init__(self, **kwargs: Any) -> None: images: dict[str, Image] = {} points: dict[str, Points] = {} diff --git a/src/torchio/transforms/inverse.py b/src/torchio/transforms/inverse.py index 1efd4d960..465f316f2 100644 --- a/src/torchio/transforms/inverse.py +++ b/src/torchio/transforms/inverse.py @@ -3,6 +3,7 @@ from __future__ import annotations import warnings +from collections.abc import Sequence from typing import Any from .compose import Compose @@ -13,7 +14,7 @@ def get_inverse_transform( - history: list[AppliedTransform], + history: Sequence[AppliedTransform], *, warn: bool = True, ignore_intensity: bool = False, diff --git a/tests/test_batch.py b/tests/test_batch.py index c7a6f04fc..b5fae391a 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -589,3 +589,10 @@ def test_adopt_history_shared_case(self) -> None: assert not rebuilt.has_divergent_history for subject in rebuilt.unbatch(): assert [t.name for t in subject.applied_transforms] == ["Gamma"] + + def test_uniform_applied_transforms_view_is_immutable(self) -> None: + result = tio.Gamma(log_gamma=0.2, per_instance=False)(self._batch()) + + assert isinstance(result.applied_transforms, tuple) + with pytest.raises(AttributeError): + result.applied_transforms.append("invalid") # type: ignore[attr-defined] diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index bbf2f5e31..1c9cdfc09 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -180,7 +180,7 @@ def test_fully_gated_records_no_history(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Affine(degrees=20.0, p=0.0)(batch) - assert result.applied_transforms == [] + assert not result.applied_transforms def test_fully_gated_inverse_preserves_float64(self) -> None: torch.manual_seed(0) From d169e8fb3e01566222cc186b92fe6bae1bbf792b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:25:31 +0100 Subject: [PATCH 06/11] Guard empty batch history views Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/data/batch_history.py | 4 ++++ tests/test_batch.py | 8 ++++++++ 2 files changed, 12 insertions(+) diff --git a/src/torchio/data/batch_history.py b/src/torchio/data/batch_history.py index 0bbb3cc7e..f82d42c8c 100644 --- a/src/torchio/data/batch_history.py +++ b/src/torchio/data/batch_history.py @@ -60,6 +60,8 @@ def history(self, index: int) -> tuple[Any, ...]: @property def has_divergent_history(self) -> bool: """Whether element histories differ.""" + if self.batch_size <= 1: + return False first = self._histories[0] return any( not _histories_equal(first, history) for history in self._histories[1:] @@ -72,6 +74,8 @@ def applied_transforms(self) -> tuple[Any, ...]: Raises: RuntimeError: If element histories differ. """ + if self.batch_size == 0: + return () if self.has_divergent_history: msg = ( "This batch has divergent element histories. Use `histories`" diff --git a/tests/test_batch.py b/tests/test_batch.py index b5fae391a..a133a1fc3 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -486,6 +486,14 @@ def test_from_images_empty_raises(self) -> None: with pytest.raises(ValueError, match="empty"): ImagesBatch.from_images([]) + def test_empty_history_view_is_defensive(self) -> None: + batch = object.__new__(ImagesBatch) + batch._data = torch.empty(0, 1, 1, 1, 1) + batch._histories = [] + + assert batch.has_divergent_history is False + assert batch.applied_transforms == () + def test_data_setter_non_5d_raises(self) -> None: from torchio.data.batch import ImagesBatch From a4ddd908d86e4ebb69bd87cf87be0380918ca1d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:25:56 +0100 Subject: [PATCH 07/11] Rename divergent history restack test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- tests/test_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_batch.py b/tests/test_batch.py index a133a1fc3..f8c435e8d 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -573,7 +573,7 @@ def _batch(self, batch_size: int = 4) -> SubjectsBatch: ] return SubjectsBatch.from_subjects(subjects) - def test_adopt_history_preserves_per_element(self) -> None: + def test_restack_preserves_divergent_history(self) -> None: torch.manual_seed(0) batch = self._batch() branched = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) From 866e76b1aac67dd667c6f6f870a9722b9523e1ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:26:19 +0100 Subject: [PATCH 08/11] Rename uniform history restack test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- tests/test_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_batch.py b/tests/test_batch.py index f8c435e8d..7a045ecd1 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -588,7 +588,7 @@ def test_restack_preserves_divergent_history(self) -> None: t.name for t in original.applied_transforms ] - def test_adopt_history_shared_case(self) -> None: + def test_restack_preserves_uniform_history(self) -> None: torch.manual_seed(0) batch = self._batch() transformed = tio.Gamma(log_gamma=0.3, per_instance=False)(batch) From 80f7b1aa8e92ec613793b8093cafefb3d1fe8539 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:38:26 +0100 Subject: [PATCH 09/11] Clarify invertible history requirements Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/data/invertible.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/torchio/data/invertible.py b/src/torchio/data/invertible.py index 34702ae98..2267754f4 100644 --- a/src/torchio/data/invertible.py +++ b/src/torchio/data/invertible.py @@ -14,8 +14,9 @@ class Invertible: Provides `apply_inverse_transform()` and `get_inverse_transform()` to undo recorded transforms. - Classes that inherit from this mixin must initialise - `self.applied_transforms = []` in their constructor. + Subclasses must expose an `applied_transforms` sequence. `Subject` + and `Image` use mutable lists; batch containers provide an immutable + uniform-history view backed by their per-element history store. """ applied_transforms: Sequence[Any] From e80d7720275f0d9216c2517b53cf7c01863cbffc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Tue, 21 Jul 2026 00:08:42 +0100 Subject: [PATCH 10/11] Normalize invalid rebatch element errors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/compose.py | 2 +- tests/test_one_of.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/torchio/transforms/compose.py b/src/torchio/transforms/compose.py index a7db004de..244dca1b0 100644 --- a/src/torchio/transforms/compose.py +++ b/src/torchio/transforms/compose.py @@ -316,7 +316,7 @@ def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any: try: batch = SubjectsBatch.from_subjects(subjects) - except ValueError as error: + except (TypeError, ValueError) as error: msg = ( f"Per-instance {transform_name} produced batch elements with" " different shapes or schemas, which cannot be re-stacked. Use" diff --git a/tests/test_one_of.py b/tests/test_one_of.py index 9887e26fa..24db9617d 100644 --- a/tests/test_one_of.py +++ b/tests/test_one_of.py @@ -181,3 +181,10 @@ def apply_transform(self, batch, params): tio.OneOf([_Spy()])(_make_subject()) assert seen == [False] + + +def test_rebatch_rejects_non_subject_elements() -> None: + from torchio.transforms.compose import _rebatch_with_history + + with pytest.raises(RuntimeError, match="cannot be re-stacked"): + _rebatch_with_history([object()], "OneOf") From 6c3ee6d6722f078abbee3c6aaa563e7850ca618c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Tue, 21 Jul 2026 00:39:36 +0100 Subject: [PATCH 11/11] Isolate uniform element history traces Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/transform.py | 6 ++++-- tests/test_per_instance.py | 7 ++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 175209a3b..40cc7a1eb 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -272,8 +272,10 @@ def _build_history_traces( """Build one clean optional history trace per element.""" batched_keys = params.get("_batched_keys") if batched_keys is None: - trace = self._make_applied_transform(_copy.deepcopy(params)) - return [trace] * batch_size + return [ + self._make_applied_transform(_copy.deepcopy(params)) + for _ in range(batch_size) + ] expected_size = params.get("_batch_size") if expected_size != batch_size: msg = ( diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index 1c9cdfc09..6e59430bf 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -108,13 +108,14 @@ def test_unbatch_slices_history(self) -> None: assert trace.params["log_gamma"] == batch_log_gammas[i] assert "_batched_keys" not in trace.params - def test_uniform_history_reuses_one_trace(self) -> None: + def test_uniform_history_uses_independent_traces(self) -> None: batch = _identical_batch(batch_size=4) result = tio.Gamma(log_gamma=0.2, per_instance=False)(batch) - trace = result.history(0)[-1] - assert all(history[-1] is trace for history in result.histories) + traces = [history[-1] for history in result.histories] + assert all(trace.params == traces[0].params for trace in traces) + assert all(trace is not traces[0] for trace in traces[1:]) class TestSpatialBatchSizeValidation: