Fix affine-elastic composition order

src/torchio/transforms/spatial/spatial.py

@@ -1568,13 +1568,14 @@ def _build_sampling_grid(
     )
 
     if affine_first:
-        # Affine first: map to input space, then add elastic offset.
-        input_voxels = _apply_voxel_mapping(output_coords, mapping)
-        input_voxels = input_voxels + displacement / input_spacing_t
-    else:
-        # Elastic first: deform in output space, then map to input.
+        # Resampling follows the inverse transform, so apply the operations
+        # in reverse order: undo the elastic field, then the affine mapping.
         deformed_output = output_coords + displacement / output_spacing_t
         input_voxels = _apply_voxel_mapping(deformed_output, mapping)
+    else:
+        # Undo the affine mapping first, then the elastic field.
+        input_voxels = _apply_voxel_mapping(output_coords, mapping)
+        input_voxels = input_voxels + displacement / input_spacing_t
 
     return input_voxels
 

tests/test_spatial.py

@@ -88,6 +88,34 @@ def test_affine_first_changes_result(self) -> None:
 
         assert not torch.allclose(first.t1.data, second.t1.data)
 
+    def test_affine_first_composes_inverse_grid_in_reverse_order(self) -> None:
+        identity = AffineMatrix(np.eye(4))
+        affine_matrix = np.diag([2.0, 1.0, 1.0, 1.0])
+        control_points = torch.zeros(4, 4, 4, 3)
+        control_points[..., 0] = 1.0
+        kwargs = {
+            "input_shape": (3, 3, 3),
+            "input_affine": identity,
+            "output_shape": (3, 3, 3),
+            "output_affine": identity,
+            "affine_matrix": affine_matrix,
+            "control_points": control_points,
+            "max_displacement": (1.0, 0.0, 0.0),
+            "device": torch.device("cpu"),
+        }
+
+        affine_then_elastic = _build_sampling_grid(affine_first=True, **kwargs)
+        elastic_then_affine = _build_sampling_grid(affine_first=False, **kwargs)
+
+        torch.testing.assert_close(
+            affine_then_elastic[2, 0, 0],
+            torch.tensor([1.5, 0.0, 0.0]),
+        )
+        torch.testing.assert_close(
+            elastic_then_affine[2, 0, 0],
+            torch.tensor([2.0, 0.0, 0.0]),
+        )
+
     def test_2d_suppresses_out_of_plane(self) -> None:
         data = torch.rand(1, 8, 8, 1)
         subject = tio.Subject(t1=tio.ScalarImage(data))