Repository: Project-MONAI/MONAI

Issue:
The Affine transform currently does not have an affine matrix parameter
**Is your feature request related to a problem? Please describe.**
The `Affine` transform forwards its parameters to the `AffineGrid` class.
But it currently does not provide the "affine" parameter which the `AffineGrid` class also has.

This allows to directly define a affine transformation matrix instead of individual parameters for transformations like rotation.
If computations are directly done with affine transformation matrices, this parameter is really helpful (necessary).

**Describe the solution you'd like**
The only change necessary, is to add the `affine` parameter and to forward it to the `AffineGrid` class.

```python
class Affine(Transform):
    ...
    def __init__(
        ...
        affine: Optional[NdarrayOrTensor] = None,
        ...
    ) -> None:
        ...
        self.affine_grid = AffineGrid(
            ...
            affine=affine,
            ...
        )
```

**Describe alternatives you've considered**
We patched the affine transform locally.


Discussion:
Hi @Spenhouet ,

Thanks for your feedback.
@wyli @ericspod I think this feature request can help improve the user experience, if you guys don't have other concerns, I will try to add it based on the suggestion here.

Thanks in advance.

Sure I feel it might be simpler to have a warp function instead of transform class (unless we cache the grid for performance)
Hi @wyli ,

I think maybe we can just add an arg `affine` to `Affine` transform:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/spatial/array.py#L1432
If `affine` is not None, ignore the `rotate_params`, `shear_params`, `translate_params`, `scale_params` args.
That would be a non-breaking change, what do you think?

Thanks.

I don't know what's the benefit here, there's not much benefit compared with "affine.resampler(img, affine.affinegrid(affine)[0])" unless we have additional caching capability
Hi @wyli ,

I think maybe @Spenhouet 's request is to set a fixed `affine` matrix in the transform chain, so adding an arg is easier to use?

Thanks.

Yes, I also do not see the harm of adding this parameter. @wyli or what is your worry?

It would also enable to easily perform the inverse operation in the chain.

We have a transform which currently contains these lines:

```python
projection = np.linalg.inv(target_affine) @ affine
transform = Affine(affine=projection[:3, :3], mode=self.mode)

transformed_matrix: torch.Tensor
transformed_matrix, _ = transform(img.squeeze(0)) 
```

Sure, we could also do this:

```python
projection = np.linalg.inv(target_affine) @ affine
transform = Affine(mode=self.mode)

transformed_matrix: torch.Tensor
transformed_matrix = transform.resampler(img.squeeze(0), transform.affinegrid(projection[:3, :3])[0])
```

But compared to simply adding the additional parameter, I don't think this is a great way to solve this.
We don't normally compose the array transforms. also, if we follow this pattern we'll build complicated transforms with many input arguments and if-else branches, then we lose the readability.
@wyli Maybe I misunderstood. We are not using the method to apply a static Affine in the pipeline.
We are in fact using it like mentioned above.

I get where you are coming from, but I wonder if this is the right case to focus on this.
The if-else branches already exist. The suggestion simply exposes the affine parameter to the Affine transform to not need to implement it like you suggested.

I don't think it will make the Affine transform any more complicated.

EDIT: But it will make other methods which use this Affine transform simpler. So we save on complexity there.

EDIT2: Or how would you rate the code complexity on these two code snippets:

```python
transform = Affine(affine=projection[:3, :3], mode=self.mode)
transformed_matrix, _ = transform(img.squeeze(0)) 
```

vs.

```python
transform = Affine(mode=self.mode)
transformed_matrix = transform.resampler(img.squeeze(0), transform.affinegrid(projection[:3, :3])[0])
```

EDIT3: Also the first one will also output the resulting affine while the second will not and will need additional code (further increasing code complexity and boilerplate).

Resolving patch:
diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -1435,6 +1435,7 @@ def __init__(
         shear_params: Optional[Union[Sequence[float], float]] = None,
         translate_params: Optional[Union[Sequence[float], float]] = None,
         scale_params: Optional[Union[Sequence[float], float]] = None,
+        affine: Optional[NdarrayOrTensor] = None,
         spatial_size: Optional[Union[Sequence[int], int]] = None,
         mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR,
         padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.REFLECTION,
@@ -1462,6 +1463,9 @@ def __init__(
                 pixel/voxel relative to the center of the input image. Defaults to no translation.
             scale_params: scale factor for every spatial dims. a tuple of 2 floats for 2D,
                 a tuple of 3 floats for 3D. Defaults to `1.0`.
+            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.
             spatial_size: output image spatial size.
                 if `spatial_size` and `self.spatial_size` are not defined, or smaller than 1,
                 the transform will use the spatial size of `img`.
@@ -1486,6 +1490,7 @@ def __init__(
             shear_params=shear_params,
             translate_params=translate_params,
             scale_params=scale_params,
+            affine=affine,
             device=device,
         )
         self.image_only = image_only


Test patch: