ENH: Speed up tractogram.apply_affine

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,10 @@ 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 True, attempt to apply the affine directly to ``pts``.
+        If False, or in-place application fails, a freshly allocated
+        array will be returned.
 
     Returns
     -------
@@ -80,8 +84,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/streamlines/array_sequence.py

@@ -146,6 +146,10 @@ def __init__(self, iterable=None, buffer_size=4):
 
         self.extend(iterable)
 
+    @property
+    def is_sliced_view(self):
+        return self._lengths.sum() != self._data.shape[0]
+
     @property
     def is_array_sequence(self):
         return True

nibabel/streamlines/tractogram.py

@@ -429,8 +429,12 @@ def apply_affine(self, affine, lazy=False):
         if np.all(affine == np.eye(4)):
             return self  # No transformation.
 
-        for i in range(len(self.streamlines)):
-            self.streamlines[i] = apply_affine(affine, self.streamlines[i])
+        if self.streamlines.is_sliced_view:
+            # Apply affine only on the selected streamlines.
+            for i in range(len(self.streamlines)):
+                self.streamlines[i] = apply_affine(affine, self.streamlines[i])
+        else:
+            self.streamlines._data = apply_affine(affine, self.streamlines._data, inplace=True)
 
         if self.affine_to_rasmm is not None:
             # Update the affine that brings back the streamlines to RASmm.

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)