NF: Add nibabel.affines.rescale_affine function

Author: effigies

nibabel/affines.py

@@ -323,3 +323,43 @@ def obliquity(affine):
     vs = voxel_sizes(affine)
     best_cosines = np.abs(affine[:-1, :-1] / vs).max(axis=1)
     return np.arccos(best_cosines)
+
+
+def rescale_affine(affine, shape, zooms, new_shape=None):
+    """ Return a new affine matrix with updated voxel sizes (zooms)
+
+    This function preserves the rotations and shears of the original
+    affine, as well as the RAS location of the central voxel of the
+    image.
+
+    Parameters
+    ----------
+    affine : (N, N) array-like
+        NxN transform matrix in homogeneous coordinates representing an affine
+        transformation from an (N-1)-dimensional space to an (N-1)-dimensional
+        space. An example is a 4x4 transform representing rotations and
+        translations in 3 dimensions.
+    shape : (N-1,) array-like
+        The extent of the (N-1) dimensions of the original space
+    zooms : (N-1,) array-like
+        The size of voxels of the output affine
+    new_shape : (N-1,) array-like, optional
+        The extent of the (N-1) dimensions of the space described by the
+        new affine. If ``None``, use ``shape``.
+
+    Returns
+    -------
+    affine : (N, N) array
+        A new affine transform with the specified voxel sizes
+
+    """
+    shape = np.array(shape, copy=False)
+    new_shape = np.array(new_shape if new_shape is not None else shape)
+
+    s = voxel_sizes(affine)
+    rzs_out = affine[:3, :3] * zooms / s
+
+    # Using xyz = A @ ijk, determine translation
+    centroid = apply_affine(affine, (shape - 1) // 2)
+    t_out = centroid - rzs_out @ ((new_shape - 1) // 2)
+    return from_matvec(rzs_out, t_out)

nibabel/tests/test_affines.py

@@ -7,7 +7,8 @@
 
 from ..eulerangles import euler2mat
 from ..affines import (AffineError, apply_affine, append_diag, to_matvec,
-                       from_matvec, dot_reduce, voxel_sizes, obliquity)
+                       from_matvec, dot_reduce, voxel_sizes, obliquity, rescale_affine)
+from ..orientations import aff2axcodes
 
 
 import pytest
@@ -192,3 +193,22 @@ def test_obliquity():
     assert_almost_equal(obliquity(aligned), [0.0, 0.0, 0.0])
     assert_almost_equal(obliquity(oblique) * 180 / pi,
                         [0.0810285, 5.1569949, 5.1569376])
+
+
+def test_rescale_affine():
+    rng = np.random.RandomState(20200415)
+    orig_shape = rng.randint(low=20, high=512, size=(3,))
+    orig_aff = np.eye(4)
+    orig_aff[:3, :] = rng.normal(size=(3, 4))
+    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)
+            assert aff2axcodes(new_aff) == orig_axcodes
+            if new_shape is None:
+                new_shape = tuple(orig_shape)
+            new_centroid = apply_affine(new_aff, (np.array(new_shape) - 1) // 2)
+            assert_almost_equal(new_centroid, orig_centroid)