Adding in additional placeholders and functions required for functional, array and dictionary transforms to operate while waiting for PR #5436

Signed-off-by: Ben Murray <ben.murray@gmail.com>

Author: atbenmurray

monai/transforms/meta_matrix.py

@@ -15,7 +15,10 @@
 import torch
 
 # placeholder to be replaced by MetaMatrix in Apply And Resample PR #5436
-from monai.transforms.utils import _create_rotate, _create_shear, _create_scale, _create_translate
+from monai.config import NdarrayOrTensor
+
+from monai.transforms.utils import _create_rotate, _create_shear, _create_scale, _create_translate, _create_rotate_90, \
+    _create_flip
 
 from monai.utils import TransformBackends
 
@@ -98,6 +101,179 @@ def translate(self, offsets: Union[Sequence[float], float], **extra_args):
         return MetaMatrix(matrix, extra_args)
 
 
+class Matrix:
+    def __init__(self, matrix: NdarrayOrTensor):
+        self.data = ensure_tensor(matrix)
+
+    # def __matmul__(self, other):
+    #     if isinstance(other, Matrix):
+    #         other_matrix = other.data
+    #     else:
+    #         other_matrix = other
+    #     return self.data @ other_matrix
+    #
+    # def __rmatmul__(self, other):
+    #     return other.__matmul__(self.data)
+
+
+class Grid:
+    def __init__(self, grid):
+        self.data = ensure_tensor(grid)
+
+    # def __matmul__(self, other):
+    #     raise NotImplementedError()
+
+
+class MetaMatrix:
+    def __init__(self, matrix: Union[NdarrayOrTensor, Matrix, Grid], metadata: Optional[dict] = None):
+
+        if not isinstance(matrix, (Matrix, Grid)):
+            if matrix.shape == 2:
+                if matrix.shape[0] != matrix.shape[1] or matrix.shape[0] not in (3, 4):
+                    raise ValueError(
+                        "If 'matrix' is passed a numpy ndarray/torch Tensor, it must"
+                        f" be 3x3 or 4x4 ('matrix' has has shape {matrix.shape})"
+                    )
+            matrix_ = Matrix(matrix)
+        else:
+            matrix_ = matrix
+        self.matrix = matrix_
+
+        self.metadata = metadata or {}
+
+    def __matmul__(self, other):
+        if isinstance(other, MetaMatrix):
+            other_ = other.matrix
+        else:
+            other_ = other
+        return MetaMatrix(self.matrix @ other_)
+
+    def __rmatmul__(self, other):
+        if isinstance(other, MetaMatrix):
+            other_ = other.matrix
+        else:
+            other_ = other
+        return MetaMatrix(other_ @ self.matrix)
+
+
+def matmul(
+    left: Union[MetaMatrix, Grid, Matrix, NdarrayOrTensor], right: Union[MetaMatrix, Grid, Matrix, NdarrayOrTensor]
+):
+    matrix_types = (MetaMatrix, Grid, Matrix, torch.Tensor, np.ndarray)
+
+    if not isinstance(left, matrix_types):
+        raise TypeError(f"'left' must be one of {matrix_types} but is {type(left)}")
+    if not isinstance(right, matrix_types):
+        raise TypeError(f"'second' must be one of {matrix_types} but is {type(right)}")
+
+    left_ = left.matrix if isinstance(left, MetaMatrix) else left
+    right_ = right.matrix if isinstance(right, MetaMatrix) else right
+
+    # TODO: it might be better to not return a metamatrix, unless we pass in the resulting
+    # metadata also
+    put_in_metamatrix = isinstance(left, MetaMatrix) or isinstance(right, MetaMatrix)
+
+    put_in_grid = isinstance(left, Grid) or isinstance(right, Grid)
+
+    put_in_matrix = isinstance(left, Matrix) or isinstance(right, Matrix)
+    put_in_matrix = False if put_in_grid is True else put_in_matrix
+
+    promote_to_tensor = not (isinstance(left_, np.ndarray) and isinstance(right_, np.ndarray))
+
+    left_raw = left_.data if isinstance(left_, (Matrix, Grid)) else left_
+    right_raw = right_.data if isinstance(right_, (Matrix, Grid)) else right_
+
+    if promote_to_tensor:
+        left_raw = torch.as_tensor(left_raw)
+        right_raw = torch.as_tensor(right_raw)
+
+    if isinstance(left_, Grid):
+        if isinstance(right_, Grid):
+            raise RuntimeError("Unable to matrix multiply two Grids")
+        else:
+            result = matmul_grid_matrix(left_raw, right_raw)
+    else:
+        if isinstance(right_, Grid):
+            result = matmul_matrix_grid(left_raw, right_raw)
+        else:
+            result = matmul_matrix_matrix(left_raw, right_raw)
+
+    if put_in_grid:
+        result = Grid(result)
+    elif put_in_matrix:
+        result = Matrix(result)
+
+    if put_in_metamatrix:
+        result = MetaMatrix(result)
+
+    return result
+
+
+def matmul_matrix_grid(left: NdarrayOrTensor, right: NdarrayOrTensor):
+    if not is_matrix_shaped(left):
+        raise ValueError(f"'left' should be a 2D or 3D homogenous matrix but has shape {left.shape}")
+
+    if not is_grid_shaped(right):
+        raise ValueError(
+            "'right' should be a 3D array with shape[0] == 2 or a "
+            f"4D array with shape[0] == 3 but has shape {right.shape}"
+        )
+
+    # flatten the grid to take advantage of torch batch matrix multiply
+    right_flat = right.reshape(right.shape[0], -1)
+    result_flat = left @ right_flat
+    # restore the grid shape
+    result = result_flat.reshape((-1,) + result_flat.shape[1:])
+    return result
+
+
+def matmul_grid_matrix(left: NdarrayOrTensor, right: NdarrayOrTensor):
+    if not is_grid_shaped(left):
+        raise ValueError(
+            "'left' should be a 3D array with shape[0] == 2 or a "
+            f"4D array with shape[0] == 3 but has shape {left.shape}"
+        )
+
+    if not is_matrix_shaped(right):
+        raise ValueError(f"'right' should be a 2D or 3D homogenous matrix but has shape {right.shape}")
+
+    try:
+        inv_matrix = torch.inverse(right)
+    except RuntimeError:
+        # the matrix is not invertible, so we will have to perform a slow grid to matrix operation
+        return matmul_grid_matrix_slow(left, right)
+
+    # invert the matrix and swap the arguments, taking advantage of
+    # matrix @ vector == vector_transposed @ matrix_inverse
+    return matmul_matrix_grid(inv_matrix, left)
+
+
+def matmul_grid_matrix_slow(left: NdarrayOrTensor, right: NdarrayOrTensor):
+    if not is_grid_shaped(left):
+        raise ValueError(
+            "'left' should be a 3D array with shape[0] == 2 or a "
+            f"4D array with shape[0] == 3 but has shape {left.shape}"
+        )
+
+    if not is_matrix_shaped(right):
+        raise ValueError(f"'right' should be a 2D or 3D homogenous matrix but has shape {right.shape}")
+
+    flat_left = left.reshape(left.shape[0], -1)
+    result_flat = torch.zeros_like(flat_left)
+    for i in range(flat_left.shape[1]):
+        vector = flat_left[:, i][None, :]
+        result_vector = vector @ right
+        result_flat[:, i] = result_vector[0, :]
+
+    # restore the grid shape
+    result = result_flat.reshape((-1,) + result_flat.shape[1:])
+    return result
+
+
+def matmul_matrix_matrix(left: NdarrayOrTensor, right: NdarrayOrTensor):
+    return left @ right
+
+
 # this will conflict with PR Replacement Apply and Resample #5436
 def apply_align_corners(matrix, spatial_size, factory):
     inflated_spatial_size = tuple(s + 1 for s in spatial_size)

monai/transforms/utils.py

@@ -763,6 +763,66 @@ def _create_rotate(
     raise ValueError(f"Unsupported spatial_dims: {spatial_dims}, available options are [2, 3].")
 
 
+def _create_rotate_90(
+        spatial_dims: int,
+        axis: Tuple[int, int],
+        steps: Optional[int] = 1,
+        eye_func: Callable = np.eye
+) -> NdarrayOrTensor:
+
+    values = [(1, 0, 0, 1),
+              (0, -1, 1, 0),
+              (-1, 0, 0, -1),
+              (0, 1, -1, 0)]
+
+    if spatial_dims == 2:
+        if axis != (0, 1):
+            raise ValueError(f"if 'spatial_dims' is 2, 'axis' must be (0, 1) but is {axis}")
+    elif spatial_dims == 3:
+        if axis not in ((0, 1), (0, 2), (1, 2)):
+            raise ValueError("if 'spatial_dims' is 3, 'axis' must be (0,1), (0, 2), or (1, 2) "
+                             f"but is {axis}")
+    else:
+        raise ValueError(f"'spatial_dims' must be 2 or 3 but is {spatial_dims}")
+
+    steps_ = steps % 4
+
+    affine = eye_func(spatial_dims + 1)
+
+    if spatial_dims == 2:
+        a, b = 0, 1
+    else:
+        a, b = axis
+
+    affine[a, a], affine[a, b], affine[b, a], affine[b, b] = values[steps]
+    return affine
+
+
+def _create_flip(
+        spatial_dims: int,
+        spatial_axis: Union[Sequence[int], int],
+        eye_func: Callable = np.eye
+):
+    affine = eye_func(spatial_dims + 1)
+    if isinstance(spatial_axis, int):
+        if spatial_axis < -spatial_dims or spatial_axis >= spatial_dims:
+            raise ValueError("'spatial_axis' values must be between "
+                             f"{-spatial_dims} and {spatial_dims-1} inclusive "
+                             f"('spatial_axis' is {spatial_axis})")
+        affine[spatial_axis, spatial_axis] = -1
+    else:
+        if any((s < -spatial_dims or s >= spatial_dims) for s in spatial_axis):
+            raise ValueError("'spatial_axis' values must be between "
+                             f"{-spatial_dims} and {spatial_dims-1} inclusive "
+                             f"('spatial_axis' is {spatial_axis})")
+
+        for i in range(spatial_dims):
+            if i in spatial_axis:
+                affine[i, i] = -1
+
+    return affine
+
+
 def create_shear(
     spatial_dims: int,
     coefs: Union[Sequence[float], float],