Support configurable rotation order in create_rotate (#8963)

Fixes #6029 .

### Description

`create_rotate` hard-coded 3D rotations to the intrinsic `Rx @ Ry @ Rz`
composition. This adds a `rotate_order` parameter following the
convention of `scipy.spatial.transform.Rotation.from_euler`: a string of
up to three axes from `{x, y, z}`, where lower case selects extrinsic
rotations (about the fixed world axes) and upper case selects intrinsic
rotations (about the moving body axes). The default `"XYZ"` reproduces
the previous behaviour exactly, so existing pipelines are unaffected.

The name avoids collision with the spline interpolation order already
selected via `mode`. The parameter is threaded through
`functional.rotate`, `Rotate`, `RandRotate`, `AffineGrid`,
`RandAffineGrid`, `Affine`, `RandAffine` and their dictionary variants.
Invalid sequences raise `ValueError`, and 2D inputs ignore it.

A new test module checks that the default matches the legacy matrix,
that every supported axis sequence matches scipy for both the numpy and
torch backends, that invalid sequences raise, that 2D inputs ignore the
order, and that the `Rotate` transform honours it while remaining
invertible.

### Types of changes
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] New tests added to cover the changes.
- [x] In-line docstrings updated.

---------

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: aymuos15

monai/transforms/spatial/array.py

@@ -936,6 +936,9 @@ class Rotate(InvertibleTransform, LazyTransform):
             the output data type is always ``float32``.
         lazy: a flag to indicate whether this transform should execute lazily or not.
             Defaults to False
+        rotate_order: for 3D inputs, the order in which the axes are rotated about, following the convention of
+            :py:func:`scipy.spatial.transform.Rotation.from_euler`. See
+            :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour).
     """
 
     backend = [TransformBackends.TORCH]
@@ -949,6 +952,7 @@ def __init__(
         align_corners: bool = False,
         dtype: DtypeLike | torch.dtype = torch.float32,
         lazy: bool = False,
+        rotate_order: str = "XYZ",
     ) -> None:
         LazyTransform.__init__(self, lazy=lazy)
         self.angle = angle
@@ -957,6 +961,7 @@ def __init__(
         self.padding_mode: str = padding_mode
         self.align_corners = align_corners
         self.dtype = dtype
+        self.rotate_order = rotate_order
 
     def __call__(
         self,
@@ -1009,6 +1014,7 @@ def __call__(
             _dtype,
             lazy=lazy_,
             transform_info=self.get_transform_info(),
+            rotate_order=self.rotate_order,
         )
 
     def inverse(self, data: torch.Tensor) -> torch.Tensor:
@@ -1741,6 +1747,10 @@ class AffineGrid(LazyTransform):
             dimensions + 1.
         lazy: a flag to indicate whether this transform should execute lazily or not.
             Defaults to False
+        rotate_order: for 3D inputs, the order in which the axes are rotated about when building the
+            rotation from ``rotate_params``, following the convention of
+            :py:func:`scipy.spatial.transform.Rotation.from_euler`. See
+            :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour).
     """
 
     backend = [TransformBackends.TORCH]
@@ -1756,6 +1766,7 @@ def __init__(
         align_corners: bool = False,
         affine: NdarrayOrTensor | None = None,
         lazy: bool = False,
+        rotate_order: str = "XYZ",
     ) -> None:
         LazyTransform.__init__(self, lazy=lazy)
         self.rotate_params = rotate_params
@@ -1767,6 +1778,7 @@ def __init__(
         self.dtype = _dtype if _dtype in (torch.float16, torch.float64, None) else torch.float32
         self.align_corners = align_corners
         self.affine = affine
+        self.rotate_order = rotate_order
 
     def __call__(
         self, spatial_size: Sequence[int] | None = None, grid: torch.Tensor | None = None, lazy: bool | None = None
@@ -1808,7 +1820,7 @@ def __call__(
         if self.affine is None:
             affine = torch.eye(spatial_dims + 1, device=_device)
             if self.rotate_params:
-                affine @= create_rotate(spatial_dims, self.rotate_params, device=_device, backend=_b)  # type: ignore[assignment]
+                affine @= create_rotate(spatial_dims, self.rotate_params, device=_device, backend=_b, rotate_order=self.rotate_order)  # type: ignore[assignment]
             if self.shear_params:
                 affine @= create_shear(spatial_dims, self.shear_params, device=_device, backend=_b)  # type: ignore[assignment]
             if self.translate_params:
@@ -2216,6 +2228,7 @@ def __init__(
         align_corners: bool = False,
         image_only: bool = False,
         lazy: bool = False,
+        rotate_order: str = "XYZ",
     ) -> None:
         """
         The affine transformations are applied in rotate, shear, translate, scale order.
@@ -2274,6 +2287,10 @@ def __init__(
             image_only: if True return only the image volume, otherwise return (image, affine).
             lazy: a flag to indicate whether this transform should execute lazily or not.
                 Defaults to False
+            rotate_order: for 3D inputs, the order in which the axes are rotated about when building the rotation
+                from ``rotate_params``, following the convention of
+                :py:func:`scipy.spatial.transform.Rotation.from_euler`. See
+                :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour).
         """
         LazyTransform.__init__(self, lazy=lazy)
         self.affine_grid = AffineGrid(
@@ -2286,6 +2303,7 @@ def __init__(
             align_corners=align_corners,
             device=device,
             lazy=lazy,
+            rotate_order=rotate_order,
         )
         self.image_only = image_only
         self.norm_coord = not normalized

monai/transforms/spatial/dictionary.py

@@ -917,6 +917,7 @@ def __init__(
         align_corners: bool = False,
         allow_missing_keys: bool = False,
         lazy: bool = False,
+        rotate_order: str = "XYZ",
     ) -> None:
         """
         Args:
@@ -969,6 +970,10 @@ def __init__(
             allow_missing_keys: don't raise exception if key is missing.
             lazy: a flag to indicate whether this transform should execute lazily or not.
                 Defaults to False
+            rotate_order: for 3D inputs, the order in which the axes are rotated about when building the rotation
+                from ``rotate_params``, following the convention of
+                :py:func:`scipy.spatial.transform.Rotation.from_euler`. See
+                :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour).
 
         See also:
             - :py:class:`monai.transforms.compose.MapTransform`
@@ -988,6 +993,7 @@ def __init__(
             dtype=dtype,  # type: ignore
             align_corners=align_corners,
             lazy=lazy,
+            rotate_order=rotate_order,
         )
         self.mode = ensure_tuple_rep(mode, len(self.keys))
         self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys))
@@ -1752,6 +1758,9 @@ class Rotated(MapTransform, InvertibleTransform, LazyTransform):
         allow_missing_keys: don't raise exception if key is missing.
         lazy: a flag to indicate whether this transform should execute lazily or not.
             Defaults to False
+        rotate_order: for 3D inputs, the order in which the axes are rotated about, following the convention of
+            :py:func:`scipy.spatial.transform.Rotation.from_euler`. See
+            :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour).
     """
 
     backend = Rotate.backend
@@ -1767,10 +1776,11 @@ def __init__(
         dtype: Sequence[DtypeLike | torch.dtype] | DtypeLike | torch.dtype = np.float32,
         allow_missing_keys: bool = False,
         lazy: bool = False,
+        rotate_order: str = "XYZ",
     ) -> None:
         MapTransform.__init__(self, keys, allow_missing_keys)
         LazyTransform.__init__(self, lazy=lazy)
-        self.rotator = Rotate(angle=angle, keep_size=keep_size, lazy=lazy)
+        self.rotator = Rotate(angle=angle, keep_size=keep_size, lazy=lazy, rotate_order=rotate_order)
 
         self.mode = ensure_tuple_rep(mode, len(self.keys))
         self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys))

monai/transforms/spatial/functional.py

@@ -383,7 +383,9 @@ def resize(
     return out.copy_meta_from(meta_info) if isinstance(out, MetaTensor) else out
 
 
-def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, lazy, transform_info):
+def rotate(
+    img, angle, output_shape, mode, padding_mode, align_corners, dtype, lazy, transform_info, rotate_order="XYZ"
+):
     """
     Functional implementation of rotate.
     This function operates eagerly or lazily according to
@@ -405,6 +407,9 @@ def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, l
             the output data type is always ``float32``.
         lazy: a flag that indicates whether the operation should be performed lazily or not
         transform_info: a dictionary with the relevant information pertaining to an applied transform.
+        rotate_order: the order in which the axes are rotated about for 3D inputs, following the convention of
+            :py:func:`scipy.spatial.transform.Rotation.from_euler`. See :py:func:`monai.transforms.utils.create_rotate`.
+            Defaults to ``"XYZ"`` (the legacy behaviour). Ignored for 2D inputs.
 
     """
 
@@ -413,7 +418,7 @@ def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, l
     if input_ndim not in (2, 3):
         raise ValueError(f"Unsupported image dimension: {input_ndim}, available options are [2, 3].")
     _angle = ensure_tuple_rep(angle, 1 if input_ndim == 2 else 3)
-    transform = create_rotate(input_ndim, _angle)
+    transform = create_rotate(input_ndim, _angle, rotate_order=rotate_order)
     if output_shape is None:
         corners = np.asarray(np.meshgrid(*[(0, dim) for dim in im_shape], indexing="ij")).reshape((len(im_shape), -1))
         corners = transform[:-1, :-1] @ corners  # type: ignore

monai/transforms/utils.py

@@ -864,6 +864,7 @@ def create_rotate(
     radians: Sequence[float] | float,
     device: torch.device | None = None,
     backend: str = TransformBackends.NUMPY,
+    rotate_order: str = "XYZ",
 ) -> NdarrayOrTensor:
     """
     create a 2D or 3D rotation matrix
@@ -872,19 +873,33 @@ def create_rotate(
         spatial_dims: {``2``, ``3``} spatial rank
         radians: rotation radians
             when spatial_dims == 3, the `radians` sequence corresponds to
-            rotation in the 1st, 2nd, and 3rd dim respectively.
+            rotation about the axes named by ``rotate_order``, in the order they are listed.
         device: device to compute and store the output (when the backend is "torch").
         backend: APIs to use, ``numpy`` or ``torch``.
+        rotate_order: the order in which the axes are rotated about when ``spatial_dims == 3``,
+            following the convention of :py:func:`scipy.spatial.transform.Rotation.from_euler`.
+            A string of up to three characters from ``{'x', 'y', 'z'}`` (or ``{'X', 'Y', 'Z'}``),
+            where ``radians[i]`` is the angle applied about ``rotate_order[i]``. Lower case letters
+            select extrinsic rotations (about the original fixed axes); upper case letters select
+            intrinsic rotations (about the moving, body-fixed axes). The default ``"XYZ"``
+            reproduces the legacy behaviour (intrinsic x, then y, then z). Ignored when
+            ``spatial_dims == 2``.
 
     Raises:
         ValueError: When ``radians`` is empty.
         ValueError: When ``spatial_dims`` is not one of [2, 3].
+        ValueError: When ``rotate_order`` is not a valid Euler axis sequence.
 
     """
     _backend = look_up_option(backend, TransformBackends)
     if _backend == TransformBackends.NUMPY:
         return _create_rotate(
-            spatial_dims=spatial_dims, radians=radians, sin_func=np.sin, cos_func=np.cos, eye_func=np.eye
+            spatial_dims=spatial_dims,
+            radians=radians,
+            sin_func=np.sin,
+            cos_func=np.cos,
+            eye_func=np.eye,
+            order=rotate_order,
         )
     if _backend == TransformBackends.TORCH:
         return _create_rotate(
@@ -893,16 +908,46 @@ def create_rotate(
             sin_func=lambda th: torch.sin(torch.as_tensor(th, dtype=torch.float32, device=device)),
             cos_func=lambda th: torch.cos(torch.as_tensor(th, dtype=torch.float32, device=device)),
             eye_func=lambda rank: torch.eye(rank, device=device),
+            order=rotate_order,
         )
     raise ValueError(f"backend {backend} is not supported")
 
 
+def _validate_euler_order(order: str, num_radians: int) -> None:
+    """
+    Validate a scipy-style Euler axis sequence.
+
+    Args:
+        order: the user-facing ``rotate_order`` value, a 1-3 character axis sequence.
+        num_radians: number of rotation angles the sequence must accommodate.
+
+    Raises:
+        ValueError: when ``order`` is not a valid Euler axis sequence for ``num_radians`` angles.
+    """
+    if not isinstance(order, str):
+        raise ValueError(f"`rotate_order` must be a string, got {type(order).__name__}.")
+    if not 1 <= len(order) <= 3:
+        raise ValueError(f"`rotate_order` must contain between 1 and 3 axes, got '{order}'.")
+    if not (order.islower() or order.isupper()):
+        raise ValueError(
+            f"`rotate_order` must be all lower case (extrinsic) or all upper case (intrinsic), got '{order}'."
+        )
+    lowered = order.lower()
+    if any(axis not in "xyz" for axis in lowered):
+        raise ValueError(f"`rotate_order` axes must be from 'x', 'y', 'z' (any case), got '{order}'.")
+    if any(lowered[i] == lowered[i + 1] for i in range(len(lowered) - 1)):
+        raise ValueError(f"`rotate_order` must not repeat the same axis consecutively, got '{order}'.")
+    if len(order) < num_radians:
+        raise ValueError(f"`rotate_order` '{order}' is too short for {num_radians} rotation angle(s).")
+
+
 def _create_rotate(
     spatial_dims: int,
     radians: Sequence[float] | float,
     sin_func: Callable = np.sin,
     cos_func: Callable = np.cos,
     eye_func: Callable = np.eye,
+    order: str = "XYZ",
 ) -> NdarrayOrTensor:
     radians = ensure_tuple(radians)
     if spatial_dims == 2:
@@ -915,30 +960,25 @@ def _create_rotate(
         raise ValueError("radians must be non empty.")
 
     if spatial_dims == 3:
-        affine = None
-        if len(radians) >= 1:
-            sin_, cos_ = sin_func(radians[0]), cos_func(radians[0])
-            affine = eye_func(4)
-            affine[1, 1], affine[1, 2] = cos_, -sin_
-            affine[2, 1], affine[2, 2] = sin_, cos_
-        if len(radians) >= 2:
-            sin_, cos_ = sin_func(radians[1]), cos_func(radians[1])
-            if affine is None:
-                raise ValueError("Affine should be a matrix.")
-            _affine = eye_func(4)
-            _affine[0, 0], _affine[0, 2] = cos_, sin_
-            _affine[2, 0], _affine[2, 2] = -sin_, cos_
-            affine = affine @ _affine
-        if len(radians) >= 3:
-            sin_, cos_ = sin_func(radians[2]), cos_func(radians[2])
-            if affine is None:
-                raise ValueError("Affine should be a matrix.")
-            _affine = eye_func(4)
-            _affine[0, 0], _affine[0, 1] = cos_, -sin_
-            _affine[1, 0], _affine[1, 1] = sin_, cos_
-            affine = affine @ _affine
-        if affine is None:
+        if len(radians) < 1:
             raise ValueError("radians must be non empty.")
+        _validate_euler_order(order, len(radians))
+        intrinsic = order.isupper()
+        affine = eye_func(4)
+        for axis, radian in zip(order.lower(), radians):
+            sin_, cos_ = sin_func(radian), cos_func(radian)
+            _affine = eye_func(4)
+            if axis == "x":
+                _affine[1, 1], _affine[1, 2] = cos_, -sin_
+                _affine[2, 1], _affine[2, 2] = sin_, cos_
+            elif axis == "y":
+                _affine[0, 0], _affine[0, 2] = cos_, sin_
+                _affine[2, 0], _affine[2, 2] = -sin_, cos_
+            else:  # axis == "z"
+                _affine[0, 0], _affine[0, 1] = cos_, -sin_
+                _affine[1, 0], _affine[1, 1] = sin_, cos_
+            # intrinsic rotations post-multiply (body-fixed axes); extrinsic pre-multiply (world axes)
+            affine = affine @ _affine if intrinsic else _affine @ affine
         return affine  # type: ignore
 
     raise ValueError(f"Unsupported spatial_dims: {spatial_dims}, available options are [2, 3].")