fix: make the spatial transforms sample the way they claim to

.github/workflows/lint.yml

@@ -51,7 +51,15 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          # The floor of the supported range, matching `python_version = "3.10"` in
+          # [tool.mypy] -- not the newest interpreter. mypy applies its target version
+          # when parsing *every* file, third-party stubs included, so a runner that
+          # resolves a dependency too new for that target breaks the job with an error
+          # in someone else's .pyi. That is what numpy did: 2.3+ requires Python >=3.11
+          # and its stubs use PEP 695 `type` statements, so a 3.12 runner installed
+          # numpy 2.5 and mypy rejected numpy/__init__.pyi as 3.12-only syntax.
+          # Installing on 3.10 resolves numpy 2.2.x, whose stubs parse under the target.
+          python-version: "3.10"
           cache: pip
           cache-dependency-path: pyproject.toml
 

smauglab/transforms/gpu/spatial.py

@@ -187,9 +187,10 @@ def apply_transform_mask(
             Convert "resample" arguments to "nearest" by default.
 
         """
-        resample_method: Resample | None
-        if "resample" in flags:
-            resample_method = flags["resample"]
+        # `resample_method` was declared but only *assigned* inside the `if`, so the
+        # restore below raised UnboundLocalError whenever flags carried no "resample".
+        resample_method: Resample | None = flags.get("resample")
+        if resample_method is not None:
             flags["resample"] = Resample.get("nearest")
         output = self.apply_transform(input, params, flags, transform)
         if resample_method is not None:
@@ -293,6 +294,24 @@ def apply_transform_mask(
         return output
 
 
+def _choose_axis(batch_size: int, device: torch.device, same_on_batch: bool) -> torch.Tensor:
+    """Pick the single axis to act on, per batch element. Returns `[B]` indices.
+
+    Drawn here, from `forward`, rather than in `make_samplers`. kornia calls
+    `make_samplers` once and caches the samplers it builds, so an axis picked there was
+    fixed for the transform's lifetime -- "degrade a random axis" degraded the *same*
+    axis for a whole training run.
+    """
+    keep = torch.randint(0, 3, (1 if same_on_batch else batch_size,), device=device)
+    return keep.expand(batch_size) if same_on_batch else keep
+
+
+def _keep_one_axis(values: torch.Tensor, keep: torch.Tensor, neutral: float) -> torch.Tensor:
+    """Keep column `keep[b]` of a `[B, 3]` draw and set the other two to `neutral`."""
+    selected = torch.arange(3, device=values.device).unsqueeze(0) == keep.unsqueeze(1)
+    return torch.where(selected, values, torch.full_like(values, neutral))
+
+
 class ScaleGenerator3D(RandomGeneratorBase):
     def __init__(self, scale: tuple[float, float], one_dim: bool = False) -> None:
         super().__init__()
@@ -301,13 +320,6 @@ def __init__(self, scale: tuple[float, float], one_dim: bool = False) -> None:
 
     def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None:
         scale = _tuple_range_reader(self.scale, 3, device, dtype)
-        if self.one_dim:
-            # Pick a random dimension to apply scaling
-            dim = torch.randint(0, 3, (1,)).item()
-            for i in range(3):
-                if i != dim:
-                    scale[i, 0] = 1.0
-                    scale[i, 1] = 1.0
         self.scalex_sampler = UniformDistribution(scale[0, 0], scale[0, 1], validate_args=False)
         self.scaley_sampler = UniformDistribution(scale[1, 0], scale[1, 1], validate_args=False)
         self.scalez_sampler = UniformDistribution(scale[2, 0], scale[2, 1], validate_args=False)
@@ -322,6 +334,10 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) ->
         scalez = _adapted_rsampling((batch_size,), self.scalez_sampler, same_on_batch)
         scale = torch.stack([scalex, scaley, scalez], dim=1)
 
+        if self.one_dim:
+            # A scale of 1.0 leaves an axis at full resolution.
+            scale = _keep_one_axis(scale, _choose_axis(batch_size, scale.device, same_on_batch), 1.0)
+
         return {"scale": torch.as_tensor(scale, device=_device, dtype=_dtype)}
 
 
@@ -449,16 +465,24 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[
 
         batch_size, C, D, H, W = input.shape
 
-        # Expect params to contain 'flip' tensor of shape [B, 3] with 0/1 values
-        # flips = None
-        # if params is not None and 'flip' in params:
-        #    flips = params['flip']
+        # params["flip"] is [B, 3] of 0/1 flags over (z, y, x), produced by
+        # FlipGenerator3D. Reading it is what makes this transform random: the loop
+        # below used to recompute the same `flip_axis`-derived list for every b and
+        # ignore the sampled flags entirely, so every call flipped all configured axes
+        # identically -- three seeded calls gave byte-identical output, and the
+        # generator (including its "at least one axis" guarantee) was dead code.
+        flips = params.get("flip")
 
         out = input.clone()
-        # For each batch element, build list of spatial dims to flip (D,H,W -> dims 2,3,4)
+        # For each batch element, build list of spatial dims to flip. `input[b]` is
+        # [C, D, H, W], so spatial axis i sits at dim 1 + i.
         for b in range(batch_size):
-            # fb expected as length-3 tensor for (z,y,x)
-            flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis]
+            if flips is None:
+                # No sampled flags (a caller invoking apply_transform directly): fall
+                # back to flipping every configured axis.
+                flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis]
+            else:
+                flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis and bool(flips[b, axis])]
 
             if len(flip_dims) > 0:
                 out[b] = torch.flip(input[b], dims=tuple(flip_dims))
@@ -649,25 +673,11 @@ def __init__(self, crop: tuple[float, float], pos: tuple[float, float], one_dim:
 
     def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None:
         crop = _tuple_range_reader(self.crop, 3, device, dtype)
-        if self.one_dim:
-            # Pick a random dimension to apply cropping
-            dim = torch.randint(0, 3, (1,)).item()
-            for i in range(3):
-                if i != dim:
-                    crop[i, 0] = 1.0
-                    crop[i, 1] = 1.0
         self.cropx_sampler = UniformDistribution(crop[0, 0], crop[0, 1], validate_args=False)
         self.cropy_sampler = UniformDistribution(crop[1, 0], crop[1, 1], validate_args=False)
         self.cropz_sampler = UniformDistribution(crop[2, 0], crop[2, 1], validate_args=False)
 
         pos = _tuple_range_reader(self.pos, 3, device, dtype)
-        if self.one_dim:
-            # Pick a random dimension to apply cropping
-            dim = torch.randint(0, 3, (1,)).item()
-            for i in range(3):
-                if i != dim:
-                    pos[i, 0] = 1.0
-                    pos[i, 1] = 1.0
         self.posx_sampler = UniformDistribution(pos[0, 0], pos[0, 1], validate_args=False)
         self.posy_sampler = UniformDistribution(pos[1, 0], pos[1, 1], validate_args=False)
         self.posz_sampler = UniformDistribution(pos[2, 0], pos[2, 1], validate_args=False)
@@ -689,4 +699,17 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) ->
         posz = _adapted_rsampling((batch_size,), self.posz_sampler, same_on_batch)
         pos = torch.stack([posx, posy, posz], dim=1)
 
+        if self.one_dim:
+            # One axis for both: `make_samplers` drew a separate `dim` for crop and for
+            # pos, so the crop could be taken along one axis while the position that
+            # placed it was randomised along another.
+            keep = _choose_axis(batch_size, crop.device, same_on_batch)
+            # A crop fraction of 1.0 keeps the whole axis. The *position*, though, is
+            # the crop centre as a fraction of the axis, so its neutral value is 0.5
+            # (centred) -- the previous code copied the crop's 1.0 onto it, which put
+            # the box centre on the far edge and left the crop flush against it after
+            # clamping.
+            crop = _keep_one_axis(crop, keep, 1.0)
+            pos = _keep_one_axis(pos, keep, 0.5)
+
         return {"crop": torch.as_tensor(crop, device=_device, dtype=_dtype), "pos": torch.as_tensor(pos, device=_device, dtype=_dtype)}

unit_tests/test_spatial_sampling.py

@@ -0,0 +1,193 @@
+"""Spatial transforms that were advertised as random but were not.
+
+The first two groups fail against the implementation that preceded them:
+
+* `RandomFlipTransformGPU` never read the flip flags its own generator sampled, so it
+  flipped every configured axis, identically, on every call.
+* The single-axis generators drew their axis in `make_samplers`, which kornia calls
+  once and caches -- so "degrade a random axis" degraded the same axis for a whole
+  training run, and `CropGenerator3D` neutralised the crop *position* to the far edge
+  instead of the centre.
+
+`TestMaskResampleRestore` is a guard rather than a regression test; see its docstring.
+"""
+
+import torch
+from kornia.constants import Resample
+
+from smauglab.transforms.gpu.base import AugmentationSequentialCustom
+from smauglab.transforms.gpu.spatial import CropGenerator3D, RandomAffine3DCustom, RandomFlipTransformGPU, ScaleGenerator3D
+from unit_tests.helpers import SmaugLabTestCase, first_output
+
+
+class TestFlipIsActuallyRandom(SmaugLabTestCase):
+    def _pipeline(self, **kwargs):
+        return AugmentationSequentialCustom(
+            RandomFlipTransformGPU(p=1.0, **kwargs),
+            data_keys=["input", "mask"],
+            same_on_batch=False,
+        )
+
+    def test_two_seeds_give_two_different_flips(self):
+        volume, seg = self.tiny_volume(), self.tiny_seg()
+        pipeline = self._pipeline(flip_axis=(0, 1, 2))
+
+        outputs = []
+        for seed in range(12):
+            torch.manual_seed(seed)
+            outputs.append(first_output(pipeline(volume.clone(), seg.clone())).clone())
+
+        distinct = {tuple(out.flatten()[:64].tolist()) for out in outputs}
+        self.assertGreater(len(distinct), 1, "every seed produced the same flip -- params['flip'] is being ignored")
+
+    def test_batch_elements_flip_independently(self):
+        """The generator samples [B, 3]; the loop used to discard it and flip all of them."""
+        torch.manual_seed(0)
+        volume = torch.rand(8, 1, 8, 8, 8)
+        seg = torch.zeros(8, 1, 8, 8, 8)
+        pipeline = self._pipeline(flip_axis=(0, 1, 2))
+
+        out = first_output(pipeline(volume.clone(), seg.clone()))
+        flipped = [bool(torch.allclose(out[b], torch.flip(volume[b], dims=(1, 2, 3)))) for b in range(8)]
+
+        self.assertIn(False, flipped, "every batch element got the identical all-axis flip")
+
+    def test_the_mask_is_flipped_the_same_way_as_the_image(self):
+        """Image and mask read the same params['flip'], so they cannot disagree."""
+        torch.manual_seed(3)
+        volume = torch.rand(4, 1, 8, 8, 8)
+        seg = (volume > 0.5).float()
+
+        result = self._pipeline(flip_axis=(0, 1, 2))(volume.clone(), seg.clone())
+        image, mask = result[0], result[1]
+
+        self.assertTrue(torch.equal((image > 0.5).float(), mask), "the mask was flipped differently from the image")
+
+    def test_only_configured_axes_are_ever_flipped(self):
+        torch.manual_seed(1)
+        volume = torch.rand(6, 1, 8, 8, 8)
+        seg = torch.zeros(6, 1, 8, 8, 8)
+
+        out = first_output(self._pipeline(flip_axis=(0,))(volume.clone(), seg.clone()))
+        for b in range(6):
+            unchanged = torch.allclose(out[b], volume[b])
+            flipped_axis0 = torch.allclose(out[b], torch.flip(volume[b], dims=(1,)))
+            self.assertTrue(unchanged or flipped_axis0, f"batch element {b} was flipped along an axis that was not configured")
+
+    def test_apply_transform_without_params_still_flips_every_configured_axis(self):
+        """The fallback for callers that reach past the generator into apply_transform."""
+        transform = RandomFlipTransformGPU(p=1.0, flip_axis=(0, 1, 2))
+        volume = torch.rand(2, 1, 8, 8, 8)
+
+        out = transform.apply_transform(volume.clone(), {}, {}, transform=None)
+
+        self.assertTrue(torch.equal(out, torch.flip(volume, dims=(2, 3, 4))))
+
+
+class TestSingleAxisIsRedrawnEveryCall(SmaugLabTestCase):
+    """The axis used to be chosen in make_samplers, which kornia calls once."""
+
+    def test_the_scale_generator_does_not_pin_one_axis_forever(self):
+        torch.manual_seed(0)
+        generator = ScaleGenerator3D(scale=(0.3, 1.0), one_dim=True)
+        generator.make_samplers(torch.device("cpu"), torch.float32)
+
+        degraded_axes = set()
+        for _ in range(20):
+            scale = generator((4,), same_on_batch=False)["scale"]
+            # Exactly one axis per row is scaled; the rest sit at the neutral 1.0.
+            for row in scale:
+                self.assertEqual(int((row != 1.0).sum()), 1, "more than one axis was degraded")
+                degraded_axes.add(int((row != 1.0).nonzero().item()))
+
+        self.assertGreater(len(degraded_axes), 1, "the same axis was degraded on every call")
+
+    def test_same_on_batch_degrades_one_shared_axis(self):
+        torch.manual_seed(0)
+        generator = ScaleGenerator3D(scale=(0.3, 1.0), one_dim=True)
+        generator.make_samplers(torch.device("cpu"), torch.float32)
+
+        scale = generator((5,), same_on_batch=True)["scale"]
+        chosen = {int((row != 1.0).nonzero().item()) for row in scale}
+        self.assertEqual(len(chosen), 1, "same_on_batch should pick one axis for the whole batch")
+
+    def test_isotropic_scaling_is_left_alone(self):
+        """one_dim=False must keep every axis independent, as before."""
+        torch.manual_seed(0)
+        generator = ScaleGenerator3D(scale=(0.3, 0.9), one_dim=False)
+        generator.make_samplers(torch.device("cpu"), torch.float32)
+
+        scale = generator((4,), same_on_batch=False)["scale"]
+        self.assertTrue(bool((scale != 1.0).all()), "one_dim=False should not neutralise any axis")
+
+    def test_the_crop_generator_uses_one_axis_for_crop_and_position(self):
+        """make_samplers drew a separate axis for each, so they could disagree."""
+        torch.manual_seed(0)
+        generator = CropGenerator3D(crop=(0.5, 0.9), pos=(0.2, 0.8), one_dim=True)
+        generator.make_samplers(torch.device("cpu"), torch.float32)
+
+        params = generator((4,), same_on_batch=False)
+        crop, pos = params["crop"], params["pos"]
+
+        for b in range(4):
+            cropped = (crop[b] != 1.0).nonzero().flatten().tolist()
+            positioned = (pos[b] != 0.5).nonzero().flatten().tolist()
+            self.assertEqual(len(cropped), 1)
+            self.assertEqual(cropped, positioned, "the crop and its position were placed on different axes")
+
+    def test_the_crop_generator_neutralises_position_at_the_centre(self):
+        """The old code copied the crop's neutral 1.0 onto `pos`, i.e. the far edge."""
+        torch.manual_seed(0)
+        generator = CropGenerator3D(crop=(0.5, 0.9), pos=(0.2, 0.8), one_dim=True)
+        generator.make_samplers(torch.device("cpu"), torch.float32)
+
+        params = generator((4,), same_on_batch=False)
+        crop, pos = params["crop"], params["pos"]
+
+        for b in range(4):
+            kept = (crop[b] != 1.0).nonzero().flatten().tolist()
+            untouched = [axis for axis in range(3) if axis != kept[0]]
+            for axis in untouched:
+                self.assertAlmostEqual(float(pos[b, axis]), 0.5, places=5, msg="a non-cropped axis was not centred")
+
+
+class TestMaskResampleRestore(SmaugLabTestCase):
+    """Guard tests for `RandomAffine3DCustom.apply_transform_mask`.
+
+    Unlike the rest of this file these pass before the change too: `resample_method`
+    was annotated but assigned only inside the `if`, so the restore below it could
+    read an unbound local -- except that `apply_transform` indexes `flags["resample"]`
+    unguarded and raises `KeyError` first, which makes the unbound local unreachable
+    rather than harmless. `flags.get(...)` removes the hazard; these tests pin the
+    behaviour that has to survive it.
+    """
+
+    def _transform(self):
+        return RandomAffine3DCustom(p=1.0, degrees=5, align_corners=True)
+
+    def _flags(self, resample: str = "bilinear"):
+        transform = self._transform()
+        flags = dict(transform.flags)
+        flags["resample"] = Resample.get(resample)
+        return transform, flags
+
+    def test_a_mask_is_resampled_and_keeps_its_shape(self):
+        transform, flags = self._flags()
+        seg = self.tiny_seg()
+        params = transform.forward_parameters(seg.shape)
+        matrix = transform.compute_transformation(seg, params, flags)
+
+        out = transform.apply_transform_mask(seg.clone(), params, flags, transform=matrix)
+
+        self.assertEqual(out.shape, seg.shape)
+
+    def test_the_callers_resample_mode_is_restored(self):
+        """The method flips the flag to "nearest" for the mask and must put it back."""
+        transform, flags = self._flags()
+        seg = self.tiny_seg()
+        params = transform.forward_parameters(seg.shape)
+        matrix = transform.compute_transformation(seg, params, flags)
+
+        transform.apply_transform_mask(seg.clone(), params, flags, transform=matrix)
+
+        self.assertEqual(flags["resample"], Resample.get("bilinear"))