NF - add keyword arguments to sform/qform get/set

Add 'coded' argument to getters; returns both affine (or None if code is
0) and code.

Add ability to pass None as affine (defaulting to code being 0)

Add `strip_shears` keyword to qform to point out shears are being
stripped.

Author: matthew-brett

nibabel/nifti1.py

@@ -701,9 +701,29 @@ def get_qform_quaternion(self):
         # Adjust threshold to fact that source data was float32
         return fillpositive(bcd, FLOAT32_EPS_3)
 
-    def get_qform(self):
-        ''' Return 4x4 affine matrix from qform parameters in header '''
+    def get_qform(self, coded=False):
+        """ Return 4x4 affine matrix from qform parameters in header
+
+        Parameters
+        ----------
+        coded : bool, optional
+            If True, return {affine or None}, and qform code.  If False, just
+            return affine.  {affine or None} means, return None if qform code ==
+            0, and affine otherwise.
+
+        Returns
+        -------
+        affine : None or (4,4) ndarray
+            If `coded` is False, always return affine reconstructed from qform
+            quaternion.  If `coded` is True, return None if qform code is 0,
+            else return the affine.
+        code : int
+            Qform code. Only returned if `coded` is True.
+        """
         hdr = self._structarr
+        code = hdr['qform_code']
+        if code == 0 and coded:
+            return None, 0
         quat = self.get_qform_quaternion()
         R = quat2mat(quat)
         vox = hdr['pixdim'][1:4].copy()
@@ -718,30 +738,39 @@ def get_qform(self):
         out = np.eye(4)
         out[0:3, 0:3] = M
         out[0:3, 3] = [hdr['qoffset_x'], hdr['qoffset_y'], hdr['qoffset_z']]
+        if coded:
+            return out, code
         return out
 
-    def set_qform(self, affine, code=None):
+    def set_qform(self, affine, code=None, strip_shears=True):
         ''' Set qform header values from 4x4 affine
 
         Parameters
         ----------
         hdr : nifti1 header
-        affine : 4x4 array
-            affine transform to write into qform
+        affine : None or 4x4 array
+            affine transform to write into sform. If None, only set code.
         code : None, string or integer
             String or integer giving meaning of transform in *affine*.
-            The default is None.  If code is None, then {if current
-            qform code is not 0, leave code as it is in the header; else
-            set to 2 ('aligned')}.
+            The default is None.  If code is None, then:
+            * If affine is None, `code`-> 0
+            * If affine not None and sform code in header == 0, `code`-> 2
+              (aligned)
+            * If affine not None and sform code in header != 0, `code`-> sform
+              code in header
+        strip_shears : bool, optional
+            Whether to strip shears in `affine`.  If True, shears will be
+            silently stripped. If False, the presence of shears will raise a
+            ``HeaderDataError``
 
         Notes
         -----
         The qform transform only encodes translations, rotations and
-        zooms. If there are shear components to the `affine` transform,
-        the written qform gives the closest approximation where the
-        rotation matrix is orthogonal. This is to allow quaternion
-        representation. The orthogonal representation enforces
-        orthogonal axes.
+        zooms. If there are shear components to the `affine` transform, and
+        `strip_shears` is True (the default), the written qform gives the
+        closest approximation where the rotation matrix is orthogonal. This is
+        to allow quaternion representation. The orthogonal representation
+        enforces orthogonal axes.
 
         Examples
         --------
@@ -765,15 +794,25 @@ def set_qform(self, affine, code=None):
         >>> hdr.set_qform(affine, code='scanner')
         >>> int(hdr['qform_code'])
         1
+        >>> hdr.set_qform(None)
+        >>> int(hdr['qform_code'])
+        0
         '''
         hdr = self._structarr
+        old_code = hdr['qform_code']
         if code is None:
-            code = hdr['qform_code']
-            if code == 0: # default is 'aligned'
-                hdr['qform_code'] = 2
-        else:
+            if affine is None:
+                code = 0
+            elif old_code == 0:
+                code = 2 # aligned
+            else:
+                code = old_code
+        else: # code set
             code = self._field_recoders['qform_code'][code]
-            hdr['qform_code'] = code
+        hdr['qform_code'] = code
+        if affine is None:
+            return
+        affine = np.asarray(affine)
         if not affine.shape == (4, 4):
             raise TypeError('Need 4x4 affine as input')
         trans = affine[:3, 3]
@@ -793,6 +832,9 @@ def set_qform(self, affine, code=None):
         # orthogonal matrix PR, to input R
         P, S, Qs = npl.svd(R)
         PR = np.dot(P, Qs)
+        if not strip_shears and not np.allclose(PR, R):
+            raise HeaderDataError("Shears in affine and `strip_shears` is "
+                                  "False")
         # Convert to quaternion
         quat = mat2quat(PR)
         # Set into header
@@ -801,13 +843,35 @@ def set_qform(self, affine, code=None):
         hdr['pixdim'][1:4] = zooms
         hdr['quatern_b'], hdr['quatern_c'], hdr['quatern_d'] = quat[1:]
 
-    def get_sform(self):
-        ''' Return sform 4x4 affine matrix from header '''
+    def get_sform(self, coded=False):
+        """ Return 4x4 affine matrix from sform parameters in header
+
+        Parameters
+        ----------
+        coded : bool, optional
+            If True, return {affine or None}, and sform code.  If False, just
+            return affine.  {affine or None} means, return None if sform code ==
+            0, and affine otherwise.
+
+        Returns
+        -------
+        affine : None or (4,4) ndarray
+            If `coded` is False, always return affine from sform fields. If
+            `coded` is True, return None if sform code is 0, else return the
+            affine.
+        code : int
+            Sform code. Only returned if `coded` is True.
+        """
         hdr = self._structarr
+        code = hdr['sform_code']
+        if code == 0 and coded:
+            return None, 0
         out = np.eye(4)
         out[0, :] = hdr['srow_x'][:]
         out[1, :] = hdr['srow_y'][:]
         out[2, :] = hdr['srow_z'][:]
+        if coded:
+            return out, code
         return out
 
     def set_sform(self, affine, code=None):
@@ -816,13 +880,16 @@ def set_sform(self, affine, code=None):
         Parameters
         ----------
         hdr : nifti1 header
-        affine : 4x4 array
-            affine transform to write into sform
+        affine : None or 4x4 array
+            affine transform to write into sform.  If None, only set `code`
         code : None, string or integer
             String or integer giving meaning of transform in *affine*.
-            The default is None.  If code is None, then {if current
-            sform code is not 0, leave code as it is in the header; else
-            set to 2 ('aligned')}.
+            The default is None.  If code is None, then:
+            * If affine is None, `code`-> 0
+            * If affine not None and sform code in header == 0, `code`-> 2
+              (aligned)
+            * If affine not None and sform code in header != 0, `code`-> sform
+              code in header
 
         Examples
         --------
@@ -846,15 +913,25 @@ def set_sform(self, affine, code=None):
         >>> hdr.set_sform(affine, code='scanner')
         >>> int(hdr['sform_code'])
         1
+        >>> hdr.set_sform(None)
+        >>> int(hdr['sform_code'])
+        0
         '''
         hdr = self._structarr
+        old_code = hdr['sform_code']
         if code is None:
-            code = hdr['sform_code']
-            if code == 0:
-                hdr['sform_code'] = 2 # aligned
-        else:
+            if affine is None:
+                code = 0
+            elif old_code == 0:
+                code = 2 # aligned
+            else:
+                code = old_code
+        else: # code set
             code = self._field_recoders['sform_code'][code]
-            hdr['sform_code'] = code
+        hdr['sform_code'] = code
+        if affine is None:
+            return
+        affine = np.asarray(affine)
         hdr['srow_x'][:] = affine[0, :]
         hdr['srow_y'][:] = affine[1, :]
         hdr['srow_z'][:] = affine[2, :]

nibabel/tests/test_nifti1.py

@@ -17,6 +17,7 @@
 from ..casting import type_info
 from ..tmpdirs import InTemporaryDirectory
 from ..spatialimages import HeaderDataError
+from ..affines import from_matvec
 from .. import nifti1 as nifti1
 from ..nifti1 import (load, Nifti1Header, Nifti1PairHeader, Nifti1Image,
                       Nifti1Pair, Nifti1Extension, Nifti1Extensions,
@@ -213,6 +214,81 @@ def test_freesurfer_hack(self):
                 assert_equal(hdr.get_data_shape(), shape)
 
 
+    def test_qform_sform(self):
+        HC = self.header_class
+        hdr = HC()
+        assert_array_equal(hdr.get_qform(), np.eye(4))
+        empty_sform = np.zeros((4,4))
+        empty_sform[-1,-1] = 1
+        assert_array_equal(hdr.get_sform(), empty_sform)
+        assert_equal(hdr.get_qform(coded=True), (None, 0))
+        assert_equal(hdr.get_sform(coded=True), (None, 0))
+        # Affine with no shears
+        nice_aff = np.diag([2, 3, 4, 1])
+        # Affine with shears
+        nasty_aff = from_matvec(np.arange(9).reshape((3,3)), [9, 10, 11])
+        fixed_aff = unshear_44(nasty_aff)
+        for in_meth, out_meth in ((hdr.set_qform, hdr.get_qform),
+                                  (hdr.set_sform, hdr.get_sform)):
+            in_meth(nice_aff, 2)
+            aff, code = out_meth(coded=True)
+            assert_array_equal(aff, nice_aff)
+            assert_equal(code, 2)
+            assert_array_equal(out_meth(), nice_aff) # non coded
+            # Affine can also be passed if code == 0, affine will be suitably set
+            in_meth(nice_aff, 0)
+            assert_equal(out_meth(coded=True), (None, 0))
+            assert_array_almost_equal(out_meth(), nice_aff)
+            # Default qform code when previous == 0 is 2
+            in_meth(nice_aff)
+            aff, code = out_meth(coded=True)
+            assert_equal(code, 2)
+            # Unless code was non-zero before
+            in_meth(nice_aff, 1)
+            in_meth(nice_aff)
+            aff, code = out_meth(coded=True)
+            assert_equal(code, 1)
+            # Can set code without modifying affine, by passing affine=None
+            assert_array_equal(aff, nice_aff) # affine same as before
+            in_meth(None, 3)
+            aff, code = out_meth(coded=True)
+            assert_array_equal(aff, nice_aff) # affine same as before
+            assert_equal(code, 3)
+            # affine is None on its own, or with code==0, resets code to 0
+            in_meth(None, 0)
+            assert_equal(out_meth(coded=True), (None, 0))
+            in_meth(None)
+            assert_equal(out_meth(coded=True), (None, 0))
+            # List works as input
+            in_meth(nice_aff.tolist())
+            assert_array_equal(out_meth(), nice_aff)
+        # Qform specifics
+        # inexact set (with shears) is OK
+        hdr.set_qform(nasty_aff, 1)
+        assert_array_almost_equal(hdr.get_qform(), fixed_aff)
+        # Unless allow_shears is False
+        assert_raises(HeaderDataError, hdr.set_qform, nasty_aff, 1, False)
+        # Reset sform, give qform a code, to test sform
+        hdr.set_sform(None)
+        hdr.set_qform(nice_aff, 1)
+        # Check sform unchanged by setting qform
+        assert_equal(hdr.get_sform(coded=True), (None, 0))
+        # Setting does change the sform ouput
+        hdr.set_sform(nasty_aff, 1)
+        aff, code = hdr.get_sform(coded=True)
+        assert_array_equal(aff, nasty_aff)
+        assert_equal(code, 1)
+
+
+def unshear_44(affine):
+    RZS = affine[:3, :3]
+    zooms = np.sqrt(np.sum(RZS * RZS, axis=0))
+    R = RZS / zooms
+    P, S, Qs = np.linalg.svd(R)
+    PR = np.dot(P, Qs)
+    return from_matvec(PR * zooms, affine[:3,3])
+
+
 class TestNifti1SingleHeader(TestNifti1PairHeader):
 
     header_class = Nifti1Header