From f22f7c00f0ca26dc7ed15bc7f39cf0a927ee8adb Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Tue, 16 Mar 2021 12:11:25 +0000
Subject: [PATCH 1/6] inverse RandAffined

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/array.py      | 73 ++++++++++++++++++--------
 monai/transforms/spatial/dictionary.py | 33 ++++++++++--
 tests/test_inverse.py                  | 19 +++++++
 tests/test_rand_affined.py             |  2 +
 4 files changed, 102 insertions(+), 25 deletions(-)

diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
index 33b8da3ebb..a916880fc7 100644
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -925,6 +925,9 @@ class AffineGrid(Transform):
         as_tensor_output: whether to output tensor instead of numpy array.
             defaults to True.
         device: device to store the output grid data.
+        affine: If applied, ignore the params (`rotate_params`, etc.) and use the
+            supplied matrix. Should be square with each side = num of image spatial
+            dimensions + 1.
 
     """
 
@@ -936,6 +939,7 @@ def __init__(
         scale_params: Optional[Union[Sequence[float], float]] = None,
         as_tensor_output: bool = True,
         device: Optional[torch.device] = None,
+        affine: Optional[Union[np.ndarray, torch.Tensor]] = None,
     ) -> None:
         self.rotate_params = rotate_params
         self.shear_params = shear_params
@@ -945,13 +949,19 @@ def __init__(
         self.as_tensor_output = as_tensor_output
         self.device = device
 
+        self.affine = affine
+
     def __call__(
-        self, spatial_size: Optional[Sequence[int]] = None, grid: Optional[Union[np.ndarray, torch.Tensor]] = None
-    ) -> Union[np.ndarray, torch.Tensor]:
+        self,
+        spatial_size: Optional[Sequence[int]] = None,
+        grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
+        return_affine: bool = False,
+    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             spatial_size: output grid size.
             grid: grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D.
+            return_affine: boolean as to whether to return the generated affine matrix or not.
 
         Raises:
             ValueError: When ``grid=None`` and ``spatial_size=None``. Incompatible values.
@@ -963,16 +973,20 @@ def __call__(
             else:
                 raise ValueError("Incompatible values: grid=None and spatial_size=None.")
 
-        spatial_dims = len(grid.shape) - 1
-        affine = np.eye(spatial_dims + 1)
-        if self.rotate_params:
-            affine = affine @ create_rotate(spatial_dims, self.rotate_params)
-        if self.shear_params:
-            affine = affine @ create_shear(spatial_dims, self.shear_params)
-        if self.translate_params:
-            affine = affine @ create_translate(spatial_dims, self.translate_params)
-        if self.scale_params:
-            affine = affine @ create_scale(spatial_dims, self.scale_params)
+        affine: Union[np.ndarray, torch.Tensor]
+        if self.affine is None:
+            spatial_dims = len(grid.shape) - 1
+            affine = np.eye(spatial_dims + 1)
+            if self.rotate_params:
+                affine = affine @ create_rotate(spatial_dims, self.rotate_params)
+            if self.shear_params:
+                affine = affine @ create_shear(spatial_dims, self.shear_params)
+            if self.translate_params:
+                affine = affine @ create_translate(spatial_dims, self.translate_params)
+            if self.scale_params:
+                affine = affine @ create_scale(spatial_dims, self.scale_params)
+        else:
+            affine = self.affine
         affine = torch.as_tensor(np.ascontiguousarray(affine), device=self.device)
 
         grid = torch.tensor(grid) if not isinstance(grid, torch.Tensor) else grid.detach().clone()
@@ -981,9 +995,10 @@ def __call__(
         grid = (affine.float() @ grid.reshape((grid.shape[0], -1)).float()).reshape([-1] + list(grid.shape[1:]))
         if grid is None or not isinstance(grid, torch.Tensor):
             raise ValueError("Unknown grid.")
-        if self.as_tensor_output:
-            return grid
-        return np.asarray(grid.cpu().numpy())
+        output: Union[np.ndarray, torch.Tensor] = grid if self.as_tensor_output else np.asarray(grid.cpu().numpy())
+        if return_affine:
+            return output, affine
+        return output
 
 
 class RandAffineGrid(RandomizableTransform):
@@ -1053,12 +1068,16 @@ def randomize(self, data: Optional[Any] = None) -> None:
         self.scale_params = self._get_rand_param(self.scale_range, 1.0)
 
     def __call__(
-        self, spatial_size: Optional[Sequence[int]] = None, grid: Optional[Union[np.ndarray, torch.Tensor]] = None
-    ) -> Union[np.ndarray, torch.Tensor]:
+        self,
+        spatial_size: Optional[Sequence[int]] = None,
+        grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
+        return_affine: bool = False,
+    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             spatial_size: output grid size.
             grid: grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D.
+            return_affine: boolean as to whether to return the generated affine matrix or not.
 
         Returns:
             a 2D (3xHxW) or 3D (4xHxWxD) grid.
@@ -1072,7 +1091,7 @@ def __call__(
             as_tensor_output=self.as_tensor_output,
             device=self.device,
         )
-        return affine_grid(spatial_size, grid)
+        return affine_grid(spatial_size, grid, return_affine)
 
 
 class RandDeformGrid(RandomizableTransform):
@@ -1298,7 +1317,7 @@ def __call__(
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
         """
         sp_size = fall_back_tuple(spatial_size or self.spatial_size, img.shape[1:])
-        grid = self.affine_grid(spatial_size=sp_size)
+        grid: torch.Tensor = self.affine_grid(spatial_size=sp_size)  # type: ignore
         return self.resampler(
             img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode
         )
@@ -1389,7 +1408,8 @@ def __call__(
         spatial_size: Optional[Union[Sequence[int], int]] = None,
         mode: Optional[Union[GridSampleMode, str]] = None,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
-    ) -> Union[np.ndarray, torch.Tensor]:
+        return_affine: bool = False,
+    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             img: shape must be (num_channels, H, W[, D]),
@@ -1404,17 +1424,26 @@ def __call__(
             padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
                 Padding mode for outside grid values. Defaults to ``self.padding_mode``.
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+            return_affine: boolean as to whether to return the generated affine matrix or not.
         """
         self.randomize()
 
         sp_size = fall_back_tuple(spatial_size or self.spatial_size, img.shape[1:])
+        affine = np.eye(len(sp_size) + 1)
         if self._do_transform:
-            grid = self.rand_affine_grid(spatial_size=sp_size)
+            out = self.rand_affine_grid(spatial_size=sp_size)
+            if return_affine:
+                grid, affine = out
+            else:
+                grid = out
         else:
             grid = create_grid(spatial_size=sp_size)
-        return self.resampler(
+        resampled = self.resampler(
             img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode
         )
+        if return_affine:
+            return resampled, affine
+        return resampled
 
 
 class Rand2DElastic(RandomizableTransform):
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
index 170006ed2b..3ef30c7dc6 100644
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -22,11 +22,13 @@
 import torch
 
 from monai.config import DtypeLike, KeysCollection
+from monai.networks.layers import AffineTransform
 from monai.networks.layers.simplelayers import GaussianFilter
 from monai.transforms.croppad.array import CenterSpatialCrop
 from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.spatial.array import (
     Affine,
+    AffineGrid,
     Flip,
     Orientation,
     Rand2DElastic,
@@ -525,7 +527,7 @@ def __call__(
         return d
 
 
-class RandAffined(RandomizableTransform, MapTransform):
+class RandAffined(RandomizableTransform, MapTransform, InvertibleTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.RandAffine`.
     """
@@ -617,16 +619,41 @@ def __call__(
 
         sp_size = fall_back_tuple(self.rand_affine.spatial_size, data[self.keys[0]].shape[1:])
         if self._do_transform:
-            grid = self.rand_affine.rand_affine_grid(spatial_size=sp_size)
+            grid, affine = self.rand_affine.rand_affine_grid(spatial_size=sp_size, return_affine=True)
         else:
             grid = create_grid(spatial_size=sp_size)
+            affine = np.eye(len(sp_size) + 1)
 
         for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
+            self.push_transform(d, key, extra_info={"affine": affine})
             d[key] = self.rand_affine.resampler(d[key], grid, mode=mode, padding_mode=padding_mode)
         return d
 
+    def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+        d = deepcopy(dict(data))
+
+        for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
+            transform = self.get_most_recent_transform(d, key)
+            orig_size = transform[InverseKeys.ORIG_SIZE.value]
+            # Create inverse transform
+            fwd_affine = transform[InverseKeys.EXTRA_INFO.value]["affine"]
+            inv_affine = np.linalg.inv(fwd_affine)
+
+            affine_grid = AffineGrid(affine=inv_affine)
+            grid: torch.Tensor = affine_grid(orig_size)  # type: ignore
+
+            # Apply inverse transform
+            out = self.rand_affine.resampler(d[key], grid, mode, padding_mode)
+
+            # Convert to numpy
+            d[key] = out if isinstance(out, np.ndarray) else out.cpu().numpy()
+
+            # Remove the applied transform
+            self.pop_transform(d, key)
+
+        return d
+
 
-class Rand2DElasticd(RandomizableTransform, MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.Rand2DElastic`.
     """
diff --git a/tests/test_inverse.py b/tests/test_inverse.py
index 0c29ea7b08..b13803a769 100644
--- a/tests/test_inverse.py
+++ b/tests/test_inverse.py
@@ -32,6 +32,7 @@
     InvertibleTransform,
     LoadImaged,
     Orientationd,
+    RandAffined,
     RandAxisFlipd,
     RandFlipd,
     Randomizable,
@@ -286,6 +287,24 @@
     )
 )
 
+
+TESTS.append(
+    (
+        "RandAffine 3d",
+        "3D",
+        1e-1,
+        RandAffined(
+            KEYS,
+            [155, 179, 192],
+            prob=1,
+            padding_mode="zeros",
+            rotate_range=[np.pi / 6, -np.pi / 5, np.pi / 7],
+            shear_range=[(0.5, 0.5)],
+            translate_range=[10, 5, -4],
+            scale_range=[(0.8, 1.2), (0.9, 1.3)],
+        ),
+    )
+)
 TESTS_COMPOSE_X2 = [(t[0] + " Compose", t[1], t[2], Compose(Compose(t[3:]))) for t in TESTS]
 
 TESTS = TESTS + TESTS_COMPOSE_X2  # type: ignore
diff --git a/tests/test_rand_affined.py b/tests/test_rand_affined.py
index 54d71ad8f7..ae2adbe3b3 100644
--- a/tests/test_rand_affined.py
+++ b/tests/test_rand_affined.py
@@ -145,6 +145,8 @@ def test_rand_affined(self, input_param, input_data, expected_val):
         res = g(input_data)
         for key in res:
             result = res[key]
+            if "_transforms" in key:
+                continue
             expected = expected_val[key] if isinstance(expected_val, dict) else expected_val
             self.assertEqual(isinstance(result, torch.Tensor), isinstance(expected, torch.Tensor))
             if isinstance(result, torch.Tensor):

From ff143290e6c485e406a7141a5687699d6054545a Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Tue, 16 Mar 2021 12:13:04 +0000
Subject: [PATCH 2/6] undo unintentional change

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/dictionary.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
index 3ef30c7dc6..ff097bf8fc 100644
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -653,7 +653,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
 
         return d
 
-
+class Rand2DElasticd(RandomizableTransform, MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.Rand2DElastic`.
     """

From 3b985b3a9f4703c6f494924c6c03beb834d37d77 Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Tue, 16 Mar 2021 12:49:38 +0000
Subject: [PATCH 3/6] inverse Affined

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/array.py      | 13 +++++++++--
 monai/transforms/spatial/dictionary.py | 32 +++++++++++++++++++++++---
 tests/test_inverse.py                  | 16 +++++++++++++
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
index a916880fc7..3ca3f9cdc9 100644
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -1300,6 +1300,7 @@ def __call__(
         spatial_size: Optional[Union[Sequence[int], int]] = None,
         mode: Optional[Union[GridSampleMode, str]] = None,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
+        return_affine: bool = False,
     ) -> Union[np.ndarray, torch.Tensor]:
         """
         Args:
@@ -1315,12 +1316,20 @@ def __call__(
             padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
                 Padding mode for outside grid values. Defaults to ``self.padding_mode``.
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+            return_affine: boolean as to whether to return the generated affine matrix or not.
         """
         sp_size = fall_back_tuple(spatial_size or self.spatial_size, img.shape[1:])
-        grid: torch.Tensor = self.affine_grid(spatial_size=sp_size)  # type: ignore
-        return self.resampler(
+        out = self.affine_grid(spatial_size=sp_size, return_affine=return_affine)
+        if return_affine:
+            grid, affine = out
+        else:
+            grid = out
+        resampled = self.resampler(
             img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode
         )
+        if return_affine:
+            return resampled, affine
+        return resampled
 
 
 class RandAffine(RandomizableTransform):
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
index ff097bf8fc..af0f8612cd 100644
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -22,7 +22,6 @@
 import torch
 
 from monai.config import DtypeLike, KeysCollection
-from monai.networks.layers import AffineTransform
 from monai.networks.layers.simplelayers import GaussianFilter
 from monai.transforms.croppad.array import CenterSpatialCrop
 from monai.transforms.inverse import InvertibleTransform
@@ -454,7 +453,7 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda
         return d
 
 
-class Affined(RandomizableTransform, MapTransform):
+class Affined(MapTransform, InvertibleTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.Affine`.
     """
@@ -523,7 +522,33 @@ def __call__(
     ) -> Dict[Hashable, Union[np.ndarray, torch.Tensor]]:
         d = dict(data)
         for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
-            d[key] = self.affine(d[key], mode=mode, padding_mode=padding_mode)
+            orig_size = d[key].shape[1:]
+            d[key], affine = self.affine(d[key], mode=mode, padding_mode=padding_mode, return_affine=True)
+            self.push_transform(d, key, orig_size=orig_size, extra_info={"affine": affine})
+        return d
+
+    def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+        d = deepcopy(dict(data))
+
+        for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
+            transform = self.get_most_recent_transform(d, key)
+            orig_size = transform[InverseKeys.ORIG_SIZE.value]
+            # Create inverse transform
+            fwd_affine = transform[InverseKeys.EXTRA_INFO.value]["affine"]
+            inv_affine = np.linalg.inv(fwd_affine)
+
+            affine_grid = AffineGrid(affine=inv_affine)
+            grid: torch.Tensor = affine_grid(orig_size)  # type: ignore
+
+            # Apply inverse transform
+            out = self.affine.resampler(d[key], grid, mode, padding_mode)
+
+            # Convert to numpy
+            d[key] = out if isinstance(out, np.ndarray) else out.cpu().numpy()
+
+            # Remove the applied transform
+            self.pop_transform(d, key)
+
         return d
 
 
@@ -653,6 +678,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
 
         return d
 
+
 class Rand2DElasticd(RandomizableTransform, MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.Rand2DElastic`.
diff --git a/tests/test_inverse.py b/tests/test_inverse.py
index b13803a769..f3b28a5642 100644
--- a/tests/test_inverse.py
+++ b/tests/test_inverse.py
@@ -23,6 +23,7 @@
 from monai.networks.nets import UNet
 from monai.transforms import (
     AddChanneld,
+    Affined,
     BorderPadd,
     CenterSpatialCropd,
     Compose,
@@ -287,6 +288,21 @@
     )
 )
 
+TESTS.append(
+    (
+        "Affine 3d",
+        "3D",
+        1e-1,
+        Affined(
+            KEYS,
+            spatial_size=[155, 179, 192],
+            rotate_params=[np.pi / 6, -np.pi / 5, np.pi / 7],
+            shear_params=[0.5, 0.5],
+            translate_params=[10, 5, -4],
+            scale_params=[0.8, 1.3],
+        ),
+    )
+)
 
 TESTS.append(
     (

From 62bbb21a729a11955d4bc6299b153935571b4e39 Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Tue, 16 Mar 2021 12:55:31 +0000
Subject: [PATCH 4/6] code format

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
index 3ca3f9cdc9..59f9f4ddcc 100644
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -1301,7 +1301,7 @@ def __call__(
         mode: Optional[Union[GridSampleMode, str]] = None,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
         return_affine: bool = False,
-    ) -> Union[np.ndarray, torch.Tensor]:
+    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             img: shape must be (num_channels, H, W[, D]),

From a3f50e216b2389d2d2a0ae68186c782179d696e5 Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Wed, 17 Mar 2021 16:27:04 +0000
Subject: [PATCH 5/6] remove return_affine

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/array.py      | 62 ++++++++++----------------
 monai/transforms/spatial/dictionary.py |  6 ++-
 tests/test_inverse.py                  |  1 +
 3 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
index a263d0e43a..f22adecbd5 100644
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -961,13 +961,11 @@ def __call__(
         self,
         spatial_size: Optional[Sequence[int]] = None,
         grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
-        return_affine: bool = False,
     ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             spatial_size: output grid size.
             grid: grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D.
-            return_affine: boolean as to whether to return the generated affine matrix or not.
 
         Raises:
             ValueError: When ``grid=None`` and ``spatial_size=None``. Incompatible values.
@@ -979,7 +977,6 @@ def __call__(
             else:
                 raise ValueError("Incompatible values: grid=None and spatial_size=None.")
 
-        affine: Union[np.ndarray, torch.Tensor]
         if self.affine is None:
             spatial_dims = len(grid.shape) - 1
             affine = np.eye(spatial_dims + 1)
@@ -991,20 +988,21 @@ def __call__(
                 affine = affine @ create_translate(spatial_dims, self.translate_params)
             if self.scale_params:
                 affine = affine @ create_scale(spatial_dims, self.scale_params)
-        else:
-            affine = self.affine
-        affine = torch.as_tensor(np.ascontiguousarray(affine), device=self.device)
+            self.affine = affine
+
+        self.affine = torch.as_tensor(np.ascontiguousarray(self.affine), device=self.device)
 
         grid = torch.tensor(grid) if not isinstance(grid, torch.Tensor) else grid.detach().clone()
         if self.device:
             grid = grid.to(self.device)
-        grid = (affine.float() @ grid.reshape((grid.shape[0], -1)).float()).reshape([-1] + list(grid.shape[1:]))
+        grid = (self.affine.float() @ grid.reshape((grid.shape[0], -1)).float()).reshape([-1] + list(grid.shape[1:]))
         if grid is None or not isinstance(grid, torch.Tensor):
             raise ValueError("Unknown grid.")
-        output: Union[np.ndarray, torch.Tensor] = grid if self.as_tensor_output else np.asarray(grid.cpu().numpy())
-        if return_affine:
-            return output, affine
-        return output
+        return grid if self.as_tensor_output else np.asarray(grid.cpu().numpy())
+
+    def get_transformation_matrix(self) -> Optional[Union[np.ndarray, torch.Tensor]]:
+        """Get the most recently applied transformation matrix"""
+        return self.affine
 
 
 class RandAffineGrid(RandomizableTransform):
@@ -1055,6 +1053,7 @@ def __init__(
 
         self.as_tensor_output = as_tensor_output
         self.device = device
+        self.affine: Optional[Union[np.ndarray, torch.Tensor]] = None
 
     def _get_rand_param(self, param_range, add_scalar: float = 0.0):
         out_param = []
@@ -1077,13 +1076,11 @@ def __call__(
         self,
         spatial_size: Optional[Sequence[int]] = None,
         grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
-        return_affine: bool = False,
     ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
         """
         Args:
             spatial_size: output grid size.
             grid: grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D.
-            return_affine: boolean as to whether to return the generated affine matrix or not.
 
         Returns:
             a 2D (3xHxW) or 3D (4xHxWxD) grid.
@@ -1097,7 +1094,13 @@ def __call__(
             as_tensor_output=self.as_tensor_output,
             device=self.device,
         )
-        return affine_grid(spatial_size, grid, return_affine)
+        grid = affine_grid(spatial_size, grid)
+        self.affine = affine_grid.get_transformation_matrix()
+        return grid
+
+    def get_transformation_matrix(self) -> Optional[Union[np.ndarray, torch.Tensor]]:
+        """Get the most recently applied transformation matrix"""
+        return self.affine
 
 
 class RandDeformGrid(RandomizableTransform):
@@ -1306,8 +1309,7 @@ def __call__(
         spatial_size: Optional[Union[Sequence[int], int]] = None,
         mode: Optional[Union[GridSampleMode, str]] = None,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
-        return_affine: bool = False,
-    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
+    ) -> Union[np.ndarray, torch.Tensor]:
         """
         Args:
             img: shape must be (num_channels, H, W[, D]),
@@ -1322,20 +1324,12 @@ def __call__(
             padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
                 Padding mode for outside grid values. Defaults to ``self.padding_mode``.
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
-            return_affine: boolean as to whether to return the generated affine matrix or not.
         """
         sp_size = fall_back_tuple(spatial_size or self.spatial_size, img.shape[1:])
-        out = self.affine_grid(spatial_size=sp_size, return_affine=return_affine)
-        if return_affine:
-            grid, affine = out
-        else:
-            grid = out
-        resampled = self.resampler(
+        grid = self.affine_grid(spatial_size=sp_size)
+        return self.resampler(
             img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode
         )
-        if return_affine:
-            return resampled, affine
-        return resampled
 
 
 class RandAffine(RandomizableTransform):
@@ -1423,8 +1417,7 @@ def __call__(
         spatial_size: Optional[Union[Sequence[int], int]] = None,
         mode: Optional[Union[GridSampleMode, str]] = None,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
-        return_affine: bool = False,
-    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
+    ) -> Union[np.ndarray, torch.Tensor]:
         """
         Args:
             img: shape must be (num_channels, H, W[, D]),
@@ -1439,26 +1432,17 @@ def __call__(
             padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
                 Padding mode for outside grid values. Defaults to ``self.padding_mode``.
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
-            return_affine: boolean as to whether to return the generated affine matrix or not.
         """
         self.randomize()
 
         sp_size = fall_back_tuple(spatial_size or self.spatial_size, img.shape[1:])
-        affine = np.eye(len(sp_size) + 1)
         if self._do_transform:
-            out = self.rand_affine_grid(spatial_size=sp_size)
-            if return_affine:
-                grid, affine = out
-            else:
-                grid = out
+            grid = self.rand_affine_grid(spatial_size=sp_size)
         else:
             grid = create_grid(spatial_size=sp_size)
-        resampled = self.resampler(
+        return self.resampler(
             img=img, grid=grid, mode=mode or self.mode, padding_mode=padding_mode or self.padding_mode
         )
-        if return_affine:
-            return resampled, affine
-        return resampled
 
 
 class Rand2DElastic(RandomizableTransform):
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
index 560b39bf79..caa1a34e08 100644
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -572,7 +572,8 @@ def __call__(
         d = dict(data)
         for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
             orig_size = d[key].shape[1:]
-            d[key], affine = self.affine(d[key], mode=mode, padding_mode=padding_mode, return_affine=True)
+            d[key] = self.affine(d[key], mode=mode, padding_mode=padding_mode)
+            affine = self.affine.affine_grid.get_transformation_matrix()
             self.push_transform(d, key, orig_size=orig_size, extra_info={"affine": affine})
         return d
 
@@ -693,7 +694,8 @@ def __call__(
 
         sp_size = fall_back_tuple(self.rand_affine.spatial_size, data[self.keys[0]].shape[1:])
         if self._do_transform:
-            grid, affine = self.rand_affine.rand_affine_grid(spatial_size=sp_size, return_affine=True)
+            grid = self.rand_affine.rand_affine_grid(spatial_size=sp_size)
+            affine = self.rand_affine.rand_affine_grid.get_transformation_matrix()
         else:
             grid = create_grid(spatial_size=sp_size)
             affine = np.eye(len(sp_size) + 1)
diff --git a/tests/test_inverse.py b/tests/test_inverse.py
index 42954a6aea..c1225ea11c 100644
--- a/tests/test_inverse.py
+++ b/tests/test_inverse.py
@@ -400,6 +400,7 @@
         ),
     )
 )
+
 TESTS_COMPOSE_X2 = [(t[0] + " Compose", t[1], t[2], Compose(Compose(t[3:]))) for t in TESTS]
 
 TESTS = TESTS + TESTS_COMPOSE_X2  # type: ignore

From e2a1dac521442f05d3df324a38bd690156d24ce7 Mon Sep 17 00:00:00 2001
From: Richard Brown <33289025+rijobro@users.noreply.github.com>
Date: Wed, 17 Mar 2021 16:35:07 +0000
Subject: [PATCH 6/6] code format

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>
---
 monai/transforms/spatial/array.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
index f22adecbd5..de9bba8e95 100644
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -961,7 +961,7 @@ def __call__(
         self,
         spatial_size: Optional[Sequence[int]] = None,
         grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
-    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
+    ) -> Union[np.ndarray, torch.Tensor]:
         """
         Args:
             spatial_size: output grid size.
@@ -1076,7 +1076,7 @@ def __call__(
         self,
         spatial_size: Optional[Sequence[int]] = None,
         grid: Optional[Union[np.ndarray, torch.Tensor]] = None,
-    ) -> Union[np.ndarray, torch.Tensor, Tuple[Union[np.ndarray, torch.Tensor], torch.Tensor]]:
+    ) -> Union[np.ndarray, torch.Tensor]:
         """
         Args:
             spatial_size: output grid size.