Add inplace option to apply_affine.

Co-authored-by: Chris Markiewicz <[email protected]>

Author: MarcCote

nibabel/affines.py

@@ -14,7 +14,7 @@ class AffineError(ValueError):
     pass
 
 
-def apply_affine(aff, pts):
+def apply_affine(aff, pts, inplace=False):
     """ Apply affine matrix `aff` to points `pts`
 
     Returns result of application of `aff` to the *right* of `pts`.  The
@@ -39,6 +39,9 @@ def apply_affine(aff, pts):
     pts : (..., N-1) array-like
         Points, where the last dimension contains the coordinates of each
         point.  For 3D, the last dimension will be length 3.
+    inplace : bool, optional
+        If False, the affine operation is done on a copy of `pts`.
+        Otherwise, `pts` gets modified directly.
 
     Returns
     -------
@@ -80,8 +83,18 @@ def apply_affine(aff, pts):
     # rzs == rotations, zooms, shears
     rzs = aff[:-1, :-1]
     trans = aff[:-1, -1]
-    res = np.dot(pts, rzs.T) + trans[None, :]
-    return res.reshape(shape)
+
+    if inplace:
+        try:
+            np.dot(pts, rzs.T, out=pts)
+        except ValueError:
+            inplace = False
+        else:
+            pts += trans[None, :]
+    if not inplace:
+        pts = pts @ rzs.T + trans[None, :]
+
+    return pts.reshape(shape)
 
 
 def to_matvec(transform):

nibabel/tests/test_affines.py

@@ -55,6 +55,12 @@ def test_apply_affine():
     pts = pts.reshape((2, 2, 3))
     exp_res = exp_res.reshape((2, 2, 3))
     assert_array_equal(apply_affine(aff, pts), exp_res)
+
+    # Check inplace modification.
+    res = apply_affine(aff, pts, inplace=True)
+    assert_array_equal(res, exp_res)
+    assert np.shares_memory(res, pts)
+
     # That ND also works
     for N in range(2, 6):
         aff = np.eye(N)
@@ -203,7 +209,7 @@ def test_rescale_affine():
     orig_zooms = voxel_sizes(orig_aff)
     orig_axcodes = aff2axcodes(orig_aff)
     orig_centroid = apply_affine(orig_aff, (orig_shape - 1) // 2)
-    
+
     for new_shape in (None, tuple(orig_shape), (256, 256, 256), (64, 64, 40)):
         for new_zooms in ((1, 1, 1), (2, 2, 3), (0.5, 0.5, 0.5)):
             new_aff = rescale_affine(orig_aff, orig_shape, new_zooms, new_shape)