Repository: TorchIO-project/torchio

Issue:
Combined Affine/Elastic augmentations
**🚀 Feature**

A combined elastic/affine augmentation. 

**Motivation**

Currently CPU use is training bottleneck for training a relatively large vision transformer, most of the time is spent in augmentation on sitk's resampling code, it looks like resampling is getting triggered once for each of the random affine and random elastic transformations.

**Pitch**

Two additional operations RandomAffineThenElastic, and RandomElasticThenAffine (I would use the later), where the two random transformations are merged and use only a single resampling.

**Notes**

Thanks for the great package!



Discussion:
Hi, @cfhammill. It would be a nice feature. Would you be available to contribute such a transform?



Patch:
diff --git a/docs/source/transforms/augmentation.rst b/docs/source/transforms/augmentation.rst
index fcf238fc7..6a9352cd9 100644
--- a/docs/source/transforms/augmentation.rst
+++ b/docs/source/transforms/augmentation.rst
@@ -71,6 +71,13 @@ Spatial
     :show-inheritance:
 
 
+:class:`RandomAffineElasticDeformation`
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. autoclass:: RandomAffineElasticDeformation
+    :show-inheritance:
+
+
 :class:`RandomAnisotropy`
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/src/torchio/transforms/__init__.py b/src/torchio/transforms/__init__.py
index e03366ce6..8e95960ab 100644
--- a/src/torchio/transforms/__init__.py
+++ b/src/torchio/transforms/__init__.py
@@ -1,5 +1,3 @@
-# noreorder
-
 from .augmentation.composition import Compose
 
 # Augmentation
@@ -23,9 +21,11 @@
 from .augmentation.intensity import Spike
 from .augmentation.intensity import Swap
 from .augmentation.spatial import Affine
+from .augmentation.spatial import AffineElasticDeformation
 from .augmentation.spatial import ElasticDeformation
 from .augmentation.spatial import Flip
 from .augmentation.spatial import RandomAffine
+from .augmentation.spatial import RandomAffineElasticDeformation
 from .augmentation.spatial import RandomAnisotropy
 from .augmentation.spatial import RandomElasticDeformation
 from .augmentation.spatial import RandomFlip
@@ -77,6 +77,8 @@
     'RandomAnisotropy',
     'RandomElasticDeformation',
     'ElasticDeformation',
+    'RandomAffineElasticDeformation',
+    'AffineElasticDeformation',
     'RandomSwap',
     'Swap',
     'RandomBlur',
diff --git a/src/torchio/transforms/augmentation/random_transform.py b/src/torchio/transforms/augmentation/random_transform.py
index 992c696a5..e6137beb6 100644
--- a/src/torchio/transforms/augmentation/random_transform.py
+++ b/src/torchio/transforms/augmentation/random_transform.py
@@ -44,9 +44,10 @@ def _get_random_seed() -> int:
         """
         return int(torch.randint(0, 2**31, (1,)).item())
 
-    def sample_uniform_sextet(self, params: TypeSextetFloat) -> TypeTripletFloat:
+    @staticmethod
+    def sample_uniform_sextet(params: TypeSextetFloat) -> TypeTripletFloat:
         results = []
         for a, b in zip(params[::2], params[1::2]):
-            results.append(self.sample_uniform(a, b))
+            results.append(RandomTransform.sample_uniform(a, b))
         sx, sy, sz = results
         return sx, sy, sz
diff --git a/src/torchio/transforms/augmentation/spatial/__init__.py b/src/torchio/transforms/augmentation/spatial/__init__.py
index 4deb7acb3..800b2eedc 100644
--- a/src/torchio/transforms/augmentation/spatial/__init__.py
+++ b/src/torchio/transforms/augmentation/spatial/__init__.py
@@ -1,5 +1,7 @@
 from .random_affine import Affine
 from .random_affine import RandomAffine
+from .random_affine_elastic_deformation import AffineElasticDeformation
+from .random_affine_elastic_deformation import RandomAffineElasticDeformation
 from .random_anisotropy import RandomAnisotropy
 from .random_elastic_deformation import ElasticDeformation
 from .random_elastic_deformation import RandomElasticDeformation
@@ -14,4 +16,6 @@
     'RandomAnisotropy',
     'RandomElasticDeformation',
     'ElasticDeformation',
+    'RandomAffineElasticDeformation',
+    'AffineElasticDeformation',
 ]
diff --git a/src/torchio/transforms/augmentation/spatial/random_affine.py b/src/torchio/transforms/augmentation/spatial/random_affine.py
index ff3b4c821..ad3dbd882 100644
--- a/src/torchio/transforms/augmentation/spatial/random_affine.py
+++ b/src/torchio/transforms/augmentation/spatial/random_affine.py
@@ -143,25 +143,25 @@ def __init__(
         )
         self.check_shape = check_shape
 
+    @staticmethod
     def get_params(
-        self,
         scales: TypeSextetFloat,
         degrees: TypeSextetFloat,
         translation: TypeSextetFloat,
         isotropic: bool,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         scaling_params = torch.as_tensor(
-            self.sample_uniform_sextet(scales),
+            RandomTransform.sample_uniform_sextet(scales),
             dtype=torch.float64,
         )
         if isotropic:
             scaling_params.fill_(scaling_params[0])
         rotation_params = torch.as_tensor(
-            self.sample_uniform_sextet(degrees),
+            RandomTransform.sample_uniform_sextet(degrees),
             dtype=torch.float64,
         )
         translation_params = torch.as_tensor(
-            self.sample_uniform_sextet(translation),
+            RandomTransform.sample_uniform_sextet(translation),
             dtype=torch.float64,
         )
         return scaling_params, rotation_params, translation_params
@@ -356,6 +356,27 @@ def get_affine_transform(self, image):
 
         return transform
 
+    def get_default_pad_value(
+        self, tensor: torch.Tensor, sitk_image: sitk.Image
+    ) -> float:
+        default_value: float
+        if self.default_pad_value == 'minimum':
+            default_value = tensor.min().item()
+        elif self.default_pad_value == 'mean':
+            default_value = get_borders_mean(
+                sitk_image,
+                filter_otsu=False,
+            )
+        elif self.default_pad_value == 'otsu':
+            default_value = get_borders_mean(
+                sitk_image,
+                filter_otsu=True,
+            )
+        else:
+            assert isinstance(self.default_pad_value, Number)
+            default_value = float(self.default_pad_value)
+        return default_value
+
     def apply_transform(self, subject: Subject) -> Subject:
         if self.check_shape:
             subject.check_consistent_spatial_shape()
@@ -374,21 +395,7 @@ def apply_transform(self, subject: Subject) -> Subject:
                     default_value = 0
                 else:
                     interpolation = self.image_interpolation
-                    if self.default_pad_value == 'minimum':
-                        default_value = tensor.min().item()
-                    elif self.default_pad_value == 'mean':
-                        default_value = get_borders_mean(
-                            sitk_image,
-                            filter_otsu=False,
-                        )
-                    elif self.default_pad_value == 'otsu':
-                        default_value = get_borders_mean(
-                            sitk_image,
-                            filter_otsu=True,
-                        )
-                    else:
-                        assert isinstance(self.default_pad_value, Number)
-                        default_value = float(self.default_pad_value)
+                    default_value = self.get_default_pad_value(tensor, sitk_image)
                 transformed_tensor = self.apply_affine_transform(
                     sitk_image,
                     transform,
diff --git a/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py
new file mode 100644
index 000000000..589fd458f
--- /dev/null
+++ b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py
@@ -0,0 +1,229 @@
+from typing import Any
+from typing import Optional
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+
+from ....constants import INTENSITY
+from ....constants import TYPE
+from ....data.io import nib_to_sitk
+from ....data.subject import Subject
+from ...spatial_transform import SpatialTransform
+from .. import RandomTransform
+from .random_affine import Affine
+from .random_elastic_deformation import ElasticDeformation
+
+
+class RandomAffineElasticDeformation(RandomTransform, SpatialTransform):
+    r"""Apply a RandomAffine and RandomElasticDeformation simultaneously.
+
+    Optimization to use only a single SimpleITK resampling. For additional details on
+    the transformations, see :class:`~torchio.transforms.RandomAffine`
+    and :class:`~torchio.transforms.RandomElasticDeformation`
+
+    Args:
+        affine_first: Apply affine before elastic deformation.
+        affine_kwargs: See :class:`~torchio.transforms.RandomAffine` for kwargs.
+        elastic_kwargs: See :class:`~torchio.transforms.RandomElasticDeformation`
+            for kwargs.
+        **kwargs: See :class:`~torchio.transforms.Transform` for additional
+            keyword arguments.
+
+    Example:
+        >>> import torchio as tio
+        >>> image = tio.datasets.Colin27().t1
+        >>> affine_kwargs = {'scales': (0.9, 1.2), 'degrees': 15}
+        >>> elastic_kwargs = {'max_displacement': (17, 12, 2)}
+        >>> transform = tio.RandomAffineElasticDeformation(
+        ...     affine_kwargs,
+        ...     elastic_kwargs
+        ... )
+        >>> transformed = transform(image)
+
+    .. plot::
+
+        import torchio as tio
+        subject = tio.datasets.Slicer('CTChest')
+        ct = subject.CT_chest
+        elastic_kwargs = {'max_displacement': (17, 12, 2)}
+        transform = tio.RandomAffineElasticDeformation(elastic_kwargs=elastic_kwargs)
+        ct_transformed = transform(ct)
+        subject.add_image(ct_transformed, 'Transformed')
+        subject.plot()
+    """
+
+    def __init__(
+        self,
+        affine_first: bool = True,
+        affine_kwargs: Optional[dict[str, Any]] = None,
+        elastic_kwargs: Optional[dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.affine_first = affine_first
+
+        # Avoid circular imports
+        from .random_affine import RandomAffine
+        from .random_elastic_deformation import RandomElasticDeformation
+
+        self.affine_kwargs = affine_kwargs or {}
+        self.random_affine = RandomAffine(**self.affine_kwargs)
+
+        self.elastic_kwargs = elastic_kwargs or {}
+        self.random_elastic = RandomElasticDeformation(**self.elastic_kwargs)
+
+    def get_params(self):
+        affine_params = self.random_affine.get_params(
+            self.random_affine.scales,
+            self.random_affine.degrees,
+            self.random_affine.translation,
+            self.random_affine.isotropic,
+        )
+        elastic_params = self.random_elastic.get_params(
+            self.random_elastic.num_control_points,
+            self.random_elastic.max_displacement,
+            self.random_elastic.num_locked_borders,
+        )
+        return affine_params, elastic_params
+
+    def apply_transform(self, subject: Subject):
+        affine_params, elastic_params = self.get_params()
+
+        scaling_params, rotation_params, translation_params = affine_params
+        affine_params = {
+            'scales': scaling_params.tolist(),
+            'degrees': rotation_params.tolist(),
+            'translation': translation_params.tolist(),
+            'center': self.random_affine.center,
+            'default_pad_value': self.random_affine.default_pad_value,
+            'image_interpolation': self.random_affine.image_interpolation,
+            'label_interpolation': self.random_affine.label_interpolation,
+            'check_shape': self.random_affine.check_shape,
+        }
+
+        elastic_params = {
+            'control_points': elastic_params,
+            'max_displacement': self.random_elastic.max_displacement,
+            'image_interpolation': self.random_elastic.image_interpolation,
+            'label_interpolation': self.random_elastic.label_interpolation,
+        }
+
+        arguments = {
+            'affine_first': self.affine_first,
+            'affine_params': affine_params,
+            'elastic_params': elastic_params,
+        }
+
+        transform = AffineElasticDeformation(**self.add_base_args(arguments))
+        transformed = transform(subject)
+        assert isinstance(transformed, Subject)
+        return transformed
+
+
+class AffineElasticDeformation(SpatialTransform):
+    r"""Apply an Affine and ElasticDeformation simultaneously.
+
+    Optimization to use only a single SimpleITK resampling. For additional details
+    on the transformations, see :class:`~torchio.transforms.augmentation.Affine`
+    and :class:`~torchio.transforms.augmentation.ElasticDeformation`
+
+    Args:
+        affine_first: Apply affine before elastic deformation.
+        affine_kwargs: See :class:`~torchio.transforms.augmentation.RandomAffine` for kwargs.
+        elastic_kwargs: See
+            :class:`~torchio.transforms.augmentation.RandomElasticDeformation` for kwargs.
+        **kwargs: See :class:`~torchio.transforms.Transform` for additional
+            keyword arguments.
+    """
+
+    def __init__(
+        self,
+        affine_first: bool,
+        affine_params: dict[str, Any],
+        elastic_params: dict[str, Any],
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.affine_first = affine_first
+
+        self.affine_params = affine_params
+        self._affine = Affine(
+            **self.affine_params,
+            **kwargs,
+        )
+        self.elastic_params = elastic_params
+        self._elastic = ElasticDeformation(
+            **self.elastic_params,
+            **kwargs,
+        )
+
+        self.args_names = ['affine_first', 'affine_params', 'elastic_params']
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        if self._affine.check_shape:
+            subject.check_consistent_spatial_shape()
+        default_value: float
+
+        for image in self.get_images(subject):
+            affine_transform = self._affine.get_affine_transform(image)
+            transformed_tensors = []
+            for tensor in image.data:
+                sitk_image = nib_to_sitk(
+                    tensor[np.newaxis],
+                    image.affine,
+                    force_3d=True,
+                )
+                if image[TYPE] != INTENSITY:
+                    interpolation = self._affine.label_interpolation
+                    default_value = 0
+                else:
+                    interpolation = self._affine.image_interpolation
+                    default_value = self._affine.get_default_pad_value(
+                        tensor, sitk_image
+                    )
+
+                bspline_transform = self._elastic.get_bspline_transform(sitk_image)
+                self._elastic.parse_free_form_transform(
+                    bspline_transform,
+                    self._elastic.max_displacement,
+                )
+
+                # stack: LIFO
+                if self.affine_first:
+                    combined_transforms = [affine_transform, bspline_transform]
+                else:
+                    combined_transforms = [bspline_transform, affine_transform]
+                composite_transform = sitk.CompositeTransform(combined_transforms)
+
+                transformed_tensor = self.apply_composite_transform(
+                    sitk_image,
+                    composite_transform,
+                    interpolation,
+                    default_value,
+                )
+                transformed_tensors.append(transformed_tensor)
+            image.set_data(torch.stack(transformed_tensors))
+        return subject
+
+    def apply_composite_transform(
+        self,
+        sitk_image: sitk.Image,
+        transform: sitk.Transform,
+        interpolation: str,
+        default_value: float,
+    ) -> torch.Tensor:
+        floating = reference = sitk_image
+
+        resampler = sitk.ResampleImageFilter()
+        resampler.SetInterpolator(self.get_sitk_interpolator(interpolation))
+        resampler.SetReferenceImage(reference)
+        resampler.SetDefaultPixelValue(float(default_value))
+        resampler.SetOutputPixelType(sitk.sitkFloat32)
+        resampler.SetTransform(transform)
+        resampled = resampler.Execute(floating)
+
+        np_array = sitk.GetArrayFromImage(resampled)
+        np_array = np_array.transpose()  # ITK to NumPy
+        tensor = torch.as_tensor(np_array)
+        return tensor
diff --git a/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py b/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
index f60d92b30..2f221a85e 100644
--- a/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
+++ b/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
@@ -1,5 +1,4 @@
 import warnings
-from collections.abc import Sequence
 from numbers import Number
 from typing import Union
 
@@ -120,8 +119,8 @@ class RandomElasticDeformation(RandomTransform, SpatialTransform):
 
     def __init__(
         self,
-        num_control_points: Union[int, tuple[int, int, int]] = 7,
-        max_displacement: Union[float, tuple[float, float, float]] = 7.5,
+        num_control_points: Union[int, TypeTripletInt] = 7,
+        max_displacement: Union[float, TypeTripletFloat] = 7.5,
         locked_borders: int = 2,
         image_interpolation: str = 'linear',
         label_interpolation: str = 'nearest',
@@ -252,7 +251,7 @@ def get_bspline_transform(
     @staticmethod
     def parse_free_form_transform(
         transform: sitk.BSplineTransform,
-        max_displacement: Sequence[TypeTripletInt],
+        max_displacement: TypeTripletFloat,
     ) -> None:
         """Issue a warning is possible folding is detected."""
         coefficient_images = transform.GetCoefficientImages()


Test patch:
diff --git a/tests/transforms/augmentation/test_random_affine_elastic_deformation.py b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py
new file mode 100644
index 000000000..500e4244a
--- /dev/null
+++ b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py
@@ -0,0 +1,353 @@
+import pytest
+import torch
+
+import torchio as tio
+
+from ...utils import TorchioTestCase
+
+
+class TestRandomAffineElasticDeformation(TorchioTestCase):
+    """Tests for `RandomAffineElasticDeformation`."""
+
+    def setUp(self):
+        # Set image origin far from center
+        super().setUp()
+        affine = self.sample_subject.t1.affine
+        affine[:3, 3] = 1e5
+
+    def test_inputs_pta_gt_one(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(p=1.5)
+
+    def test_inputs_pta_lt_zero(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(p=-1)
+
+    def test_inputs_interpolation_int(self):
+        with pytest.raises(TypeError):
+            tio.RandomAffineElasticDeformation(image_interpolation=1)
+
+    def test_inputs_interpolation(self):
+        with pytest.raises(TypeError):
+            tio.RandomAffineElasticDeformation(image_interpolation=0)
+
+    def test_num_control_points_noint(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                elastic_kwargs={'num_control_points': 2.5}
+            )
+
+    def test_num_control_points_small(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(elastic_kwargs={'num_control_points': 3})
+
+    def test_max_displacement_no_num(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                elastic_kwargs={'max_displacement': None}
+            )
+
+    def test_max_displacement_negative(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(elastic_kwargs={'max_displacement': -1})
+
+    def test_wrong_locked_borders(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(elastic_kwargs={'locked_borders': -1})
+
+    def test_coarse_grid_removed(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                elastic_kwargs={'num_control_points': (4, 5, 6), 'locked_borders': 2}
+            )
+
+    def test_folding(self):
+        # Assume shape is (10, 20, 30) and spacing is (1, 1, 1)
+        # Then grid spacing is (10/(12-2), 20/(5-2), 30/(5-2))
+        # or (1, 6.7, 10), and half is (0.5, 3.3, 5)
+        transform = tio.RandomAffineElasticDeformation(
+            elastic_kwargs={'num_control_points': (12, 5, 5), 'max_displacement': 6}
+        )
+        with pytest.warns(RuntimeWarning):
+            transform(self.sample_subject)
+
+    def test_num_control_points(self):
+        tio.RandomAffineElasticDeformation(elastic_kwargs={'num_control_points': 5})
+        tio.RandomAffineElasticDeformation(
+            elastic_kwargs={'num_control_points': (5, 6, 7)}
+        )
+
+    def test_max_displacement(self):
+        tio.RandomAffineElasticDeformation(elastic_kwargs={'max_displacement': 5})
+        tio.RandomAffineElasticDeformation(
+            elastic_kwargs={'max_displacement': (5, 6, 7)}
+        )
+
+    def test_no_displacement(self):
+        transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 0},
+            elastic_kwargs={'max_displacement': 0},
+        )
+        transformed = transform(self.sample_subject)
+        self.assert_tensor_equal(
+            self.sample_subject.t1.data,
+            transformed.t1.data,
+        )
+        self.assert_tensor_equal(
+            self.sample_subject.label.data,
+            transformed.label.data,
+        )
+
+    def test_rotation_image(self):
+        # Rotation around image center
+        transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={
+                'degrees': (90, 90),
+                'default_pad_value': 0,
+                'center': 'image',
+            }
+        )
+        transformed = transform(self.sample_subject)
+        total = transformed.t1.data.sum()
+        self.assertNotEqual(total, 0)
+
+    def test_rotation_origin(self):
+        # Rotation around far away point, image should be empty
+        transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={
+                'degrees': (90, 90),
+                'default_pad_value': 0,
+                'center': 'origin',
+            }
+        )
+        transformed = transform(self.sample_subject)
+        total = transformed.t1.data.sum()
+        assert total == 0
+
+    def test_no_rotation(self):
+        transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={
+                'scales': (1, 1),
+                'degrees': (0, 0),
+                'default_pad_value': 0,
+                'center': 'image',
+            },
+            elastic_kwargs={'max_displacement': 0},
+        )
+        transformed = transform(self.sample_subject)
+        self.assert_tensor_almost_equal(
+            self.sample_subject.t1.data,
+            transformed.t1.data,
+        )
+        transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={
+                'scales': (1, 1),
+                'degrees': (180, 180),
+                'default_pad_value': 0,
+                'center': 'image',
+            },
+            elastic_kwargs={'max_displacement': 0},
+        )
+        transformed = transform(self.sample_subject)
+        transformed = transform(transformed)
+        self.assert_tensor_almost_equal(
+            self.sample_subject.t1.data,
+            transformed.t1.data,
+        )
+
+    def test_isotropic(self):
+        tio.RandomAffineElasticDeformation(affine_kwargs={'isotropic': True})(
+            self.sample_subject
+        )
+
+    def test_mean(self):
+        tio.RandomAffineElasticDeformation(affine_kwargs={'default_pad_value': 'mean'})(
+            self.sample_subject
+        )
+
+    def test_otsu(self):
+        tio.RandomAffineElasticDeformation(affine_kwargs={'default_pad_value': 'otsu'})(
+            self.sample_subject
+        )
+
+    def test_bad_center(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'center': 'bad'})
+
+    def test_negative_scales(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (-1, 1)})
+
+    def test_scale_too_large(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 1.5})
+
+    def test_scales_range_with_negative_min(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (-1, 4)})
+
+    def test_wrong_scales_type(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 'wrong'})
+
+    def test_wrong_degrees_type(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 'wrong'})
+
+    def test_too_many_translation_values(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                affine_kwargs={'translation': (-10, 4, 42)}
+            )
+
+    def test_wrong_translation_type(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 'wrong'})
+
+    def test_wrong_center(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'center': 0})
+
+    def test_wrong_default_pad_value(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                affine_kwargs={'default_pad_value': 'wrong'}
+            )
+
+    def test_wrong_image_interpolation_type(self):
+        with pytest.raises(TypeError):
+            tio.RandomAffineElasticDeformation(affine_kwargs={'image_interpolation': 0})
+
+    def test_wrong_image_interpolation_value(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                affine_kwargs={'image_interpolation': 'wrong'}
+            )
+
+    def test_incompatible_args_isotropic(self):
+        with pytest.raises(ValueError):
+            tio.RandomAffineElasticDeformation(
+                affine_kwargs={'scales': (0.8, 0.5, 0.1), 'isotropic': True}
+            )
+
+    def test_parse_scales(self):
+        def do_assert(transform):
+            assert transform.random_affine.scales == 3 * (0.9, 1.1)
+
+        do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 0.1}))
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (0.9, 1.1)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 3 * (0.1,)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 3 * [0.9, 1.1]})
+        )
+
+    def test_parse_degrees(self):
+        def do_assert(transform):
+            assert transform.random_affine.degrees == 3 * (-10, 10)
+
+        do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 10}))
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': (-10, 10)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 3 * (10,)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 3 * [-10, 10]})
+        )
+
+    def test_parse_translation(self):
+        def do_assert(transform):
+            assert transform.random_affine.translation == 3 * (-10, 10)
+
+        do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 10}))
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'translation': (-10, 10)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 3 * (10,)})
+        )
+        do_assert(
+            tio.RandomAffineElasticDeformation(
+                affine_kwargs={'translation': 3 * [-10, 10]}
+            )
+        )
+
+    def test_default_value_label_map(self):
+        # From https://github.com/fepegar/torchio/issues/626
+        a = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).reshape(1, 3, 3, 1)
+        image = tio.LabelMap(tensor=a)
+        aff = tio.RandomAffineElasticDeformation(
+            affine_kwargs={'translation': (0, 1, 1), 'default_pad_value': 'otsu'}
+        )
+        transformed = aff(image)
+        assert all(n in (0, 1) for n in transformed.data.flatten())
+
+    def test_no_inverse(self):
+        tensor = torch.zeros((1, 2, 2, 2))
+        tensor[0, 1, 1, 1] = 1  # most RAS voxel
+        expected = torch.zeros((1, 2, 2, 2))
+        expected[0, 0, 1, 1] = 1
+        scales = 1, 1, 1
+        degrees = 0, 0, 90  # anterior should go left
+        translation = 0, 0, 0
+        apply_affine = tio.Affine(
+            scales,
+            degrees,
+            translation,
+        )
+        transformed = apply_affine(tensor)
+        self.assert_tensor_almost_equal(transformed, expected)
+
+    def test_different_spaces(self):
+        t1 = self.sample_subject.t1
+        label = tio.Resample(2)(self.sample_subject.label)
+        new_subject = tio.Subject(t1=t1, label=label)
+        with pytest.raises(RuntimeError):
+            tio.RandomAffineElasticDeformation()(new_subject)
+        tio.RandomAffineElasticDeformation(affine_kwargs={'check_shape': False})(
+            new_subject
+        )
+
+    def test_transform_order(self):
+        src_transform = tio.RandomAffineElasticDeformation(
+            affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 1},
+            elastic_kwargs={'num_control_points': 5, 'max_displacement': 1},
+        )
+
+        (scales, degrees, translation), control_points = src_transform.get_params()
+
+        max_displacement = src_transform.random_elastic.max_displacement
+
+        transform1 = tio.AffineElasticDeformation(
+            affine_first=True,
+            affine_params={
+                'scales': scales,
+                'degrees': degrees,
+                'translation': translation,
+            },
+            elastic_params={
+                'control_points': control_points,
+                'max_displacement': max_displacement,
+            },
+        )
+        transform2 = tio.AffineElasticDeformation(
+            affine_first=False,
+            affine_params={
+                'scales': scales,
+                'degrees': degrees,
+                'translation': translation,
+            },
+            elastic_params={
+                'control_points': control_points,
+                'max_displacement': max_displacement,
+            },
+        )
+
+        transformed1 = transform1(self.sample_subject)
+        transformed2 = transform2(self.sample_subject)
+        self.assert_tensor_not_equal(transformed1.t1.data, transformed2.t1.data)
diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py
index 5fabb9f9f..4e078142e 100644
--- a/tests/transforms/test_transforms.py
+++ b/tests/transforms/test_transforms.py
@@ -18,6 +18,9 @@ def get_transform(self, channels, is_3d=True, labels=True):
         landmarks_dict = {channel: np.linspace(0, 100, 13) for channel in channels}
         disp = 1 if is_3d else (1, 1, 0.01)
         elastic = tio.RandomElasticDeformation(max_displacement=disp)
+        affine_elastic = tio.RandomAffineElasticDeformation(
+            elastic_kwargs={'max_displacement': disp}
+        )
         cp_args = (9, 21, 30) if is_3d else (21, 30, 1)
         resize_args = (10, 20, 30) if is_3d else (10, 20, 1)
         flip_axes = axes_downsample = (0, 1, 2) if is_3d else (0, 1)
@@ -47,6 +50,7 @@ def get_transform(self, channels, is_3d=True, labels=True):
             tio.HistogramStandardization(landmarks_dict),
             elastic,
             tio.RandomAffine(),
+            affine_elastic,
             tio.OneOf(
                 {
                     tio.RandomAffine(): 3,