Merge pull request #339 from effigies/fix/fixcoeff_orientations

FIX: Reorient phase-encoding directions along with fieldmaps when preparing inputs to TOPUP

Author: effigies

sdcflows/interfaces/bspline.py

@@ -41,6 +41,7 @@
 )
 
 from sdcflows.transform import grid_bspline_weights
+from sdcflows.utils.tools import reorient_pedir
 
 
 LOW_MEM_BLOCK_SIZE = 1000
@@ -518,38 +519,38 @@ def _fix_topup_fieldcoeff(in_coeff, fmap_ref, pe_dir, out_file=None):
     from pathlib import Path
     import numpy as np
     import nibabel as nb
-    from sdcflows.utils.tools import ensure_positive_cosines
 
     if out_file is None:
         out_file = Path("coefficients.nii.gz").absolute()
 
-    # Load reference image
     refnii = nb.load(fmap_ref)
-
-    # Load coefficients
     coeffnii = nb.load(in_coeff)
+
+    # Coefficients generated by TOPUP are in LAS space, and we will convert to RAS.
+    # Reorient the reference image and phase-encoding direction to RAS
+    ref_ornt = nb.io_orientation(refnii.affine)
+    refnii_ras = refnii.as_reoriented(ref_ornt)
+    coeff_pe_dir = reorient_pedir(pe_dir, ref_ornt)
+
     # Coefficients - flip LR and overwrite coeffnii variable
     # Internal data orientation of FSL is LAS, so coefficients will be LR flipped,
     # and because the affine does not encode orientation (factors instead), this flip
     # always is implicit.
     # If the PE direction is x/i, the flip in the axis direction causes that the
     # fieldmap estimation must also be inverted in direction (multiply by -1.0)
-    reverse_pe = -1.0 if "i" in pe_dir.replace("x", "i") else 1.0
-    lr_axis = "".join(nb.aff2axcodes(coeffnii.affine)).index("R")
+    reverse_pe = -1.0 if coeff_pe_dir[0] == "i" else 1.0
     coeffnii = coeffnii.__class__(
-        reverse_pe * np.flip(np.asanyarray(coeffnii.dataobj), axis=lr_axis),
+        reverse_pe * np.flip(np.asanyarray(coeffnii.dataobj), axis=0),
         coeffnii.affine,
         coeffnii.header,
     )
-    # Ensure reference has positive director cosines
-    refnii, ref_axcodes = ensure_positive_cosines(refnii)
 
     # Get matrix of B-Spline control knots
     coeff_shape = np.array(coeffnii.shape[:3])
     # Get factors (w.r.t. reference's pixel sizes) to calculate separation btw control points
     factors = np.array(coeffnii.header.get_zooms()[:3])
     # Shape checking
-    ref_shape = np.array(refnii.shape[:3])
+    ref_shape = np.array(refnii_ras.shape[:3])
     exp_shape = ref_shape // factors + 3 * (factors > 1)
     if not np.all(coeff_shape == exp_shape):
         raise ValueError(
@@ -559,8 +560,8 @@ def _fix_topup_fieldcoeff(in_coeff, fmap_ref, pe_dir, out_file=None):
 
     # Contextualize the control points in space with a proper NIfTI affine
     newaff = np.eye(4)
-    newaff[:3, :3] = refnii.affine[:3, :3] * factors
-    c_ref = nb.affines.apply_affine(refnii.affine, 0.5 * (ref_shape - 1))
+    newaff[:3, :3] = refnii_ras.affine[:3, :3] * factors
+    c_ref = nb.affines.apply_affine(refnii_ras.affine, 0.5 * (ref_shape - 1))
     c_coeff = nb.affines.apply_affine(newaff, 0.5 * (coeff_shape - 1))
     newaff[:3, 3] = c_ref - c_coeff
 
@@ -621,18 +622,24 @@ def _b0_resampler(in_file, coeffs, pe, ro, hmc_xfm=None, unwarp=None, newpath=No
     # Load distorted image
     distorted_img = nb.load(in_file)
 
-    if unwarp.fit(distorted_img):
+    # Reorient to RAS to ensure consistency with coefficients
+    # The b-spline weight matrix is sensitive to orientation
+    ornt = nb.io_orientation(distorted_img.affine)
+    distorted_ras = distorted_img.as_reoriented(ornt)
+    pe_ras = reorient_pedir(pe, ornt)
+
+    if unwarp.fit(distorted_ras):
         unwarp.mapped.to_filename(retval[2])
     else:
         retval[2] = None
 
     # Unwarp
-    unwarped_img = unwarp.apply(distorted_img, ro_time=ro, pe_dir=pe)
+    unwarped_img = unwarp.apply(distorted_ras, ro_time=ro, pe_dir=pe_ras)
 
     # Write out to disk
     unwarped_img.to_filename(retval[0])
 
     # Store the corresponding spatial transformation
-    unwarp.to_displacements(ro_time=ro, pe_dir=pe).to_filename(retval[1])
+    unwarp.to_displacements(ro_time=ro, pe_dir=pe_ras).to_filename(retval[1])
 
     return retval

sdcflows/interfaces/tests/test_bspline.py

@@ -72,7 +72,6 @@ def test_bsplines(tmp_path, testnum):
     # Approximate the interpolated target
     test2 = BSplineApprox(
         in_data=test1.outputs.out_field,
-        in_mask=str(tmp_path / "target.nii.gz"),
         bs_spacing=[(4, 6, 8)],
         zooms_min=0,
         recenter=False,

sdcflows/interfaces/utils.py

@@ -21,6 +21,7 @@
 #     https://www.nipreps.org/community/licensing/
 #
 """Utilities."""
+from itertools import product
 from nipype.interfaces.base import (
     BaseInterfaceInputSpec,
     TraitedSpec,
@@ -37,6 +38,8 @@
 )
 from nipype.interfaces.mixins import CopyHeaderInterface as _CopyHeaderInterface
 
+from sdcflows.utils.tools import reorient_pedir
+
 OBLIQUE_THRESHOLD_DEG = 0.5
 
 
@@ -147,6 +150,75 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class _ReorientImageAndMetadataInputSpec(TraitedSpec):
+    in_file = File(exists=True, mandatory=True, desc="Input 3- or 4D image")
+    target_orientation = traits.Str(
+        desc="Axis codes of coordinate system to reorient to"
+    )
+    pe_dir = InputMultiObject(
+        traits.Enum(
+            *["".join(p) for p in product("ijkxyz", ("", "-"))],
+            mandatory=True,
+            desc="Phase encoding direction",
+        )
+    )
+
+
+class _ReorientImageAndMetadataOutputSpec(TraitedSpec):
+    out_file = File(desc="Reoriented image")
+    pe_dir = OutputMultiObject(
+        traits.Enum(
+            *["".join(p) for p in product("ijkxyz", ("", "-"))],
+            desc="Phase encoding direction in reoriented image",
+        )
+    )
+
+
+class ReorientImageAndMetadata(SimpleInterface):
+    input_spec = _ReorientImageAndMetadataInputSpec
+    output_spec = _ReorientImageAndMetadataOutputSpec
+
+    def _run_interface(self, runtime):
+        import numpy as np
+        import nibabel as nb
+        from nipype.utils.filemanip import fname_presuffix
+
+        target = self.inputs.target_orientation.upper()
+        if not all(code in "RASLPI" for code in target):
+            raise ValueError(
+                f"Invalid orientation code {self.inputs.target_orientation}"
+            )
+
+        img = nb.load(self.inputs.in_file)
+        img2target = nb.orientations.ornt_transform(
+            nb.io_orientation(img.affine),
+            nb.orientations.axcodes2ornt(self.inputs.target_orientation),
+        ).astype(int)
+
+        # Identity transform
+        if np.array_equal(img2target, [[0, 1], [1, 1], [2, 1]]):
+            self._results = dict(
+                out_file=self.inputs.in_file,
+                pe_dir=self.inputs.pe_dir,
+            )
+            return runtime
+
+        reoriented = img.as_reoriented(img2target)
+
+        pe_dirs = [reorient_pedir(pe_dir, img2target) for pe_dir in self.inputs.pe_dir]
+
+        self._results = dict(
+            out_file=fname_presuffix(
+                self.inputs.in_file, suffix=target, newpath=runtime.cwd
+            ),
+            pe_dir=pe_dirs,
+        )
+
+        reoriented.to_filename(self._results["out_file"])
+
+        return runtime
+
+
 class _ConvertWarpInputSpec(BaseInterfaceInputSpec):
     in_file = File(exists=True, mandatory=True, desc="output of 3dQwarp")
 

sdcflows/utils/tools.py

@@ -70,8 +70,8 @@ def ensure_positive_cosines(img):
     """
     import nibabel as nb
 
-    img_axcodes = nb.aff2axcodes(img.affine)
-    in_ornt = nb.orientations.axcodes2ornt(img_axcodes)
+    in_ornt = nb.io_orientation(img.affine)
+    img_axcodes = nb.orientations.ornt2axcodes(in_ornt)
     out_ornt = in_ornt.copy()
     out_ornt[:, 1] = 1
     ornt_xfm = nb.orientations.ornt_transform(in_ornt, out_ornt)
@@ -154,3 +154,70 @@ def brain_masker(in_file, out_file=None, padding=5):
     img.__class__(data, img.affine, img.header).to_filename(out_brain)
 
     return str(out_brain), str(out_probseg), str(out_mask)
+
+
+def reorient_pedir(pe_dir, source_ornt, target_ornt=None):
+    """Return updated PhaseEncodingDirection to account for an image data array rotation
+
+    Orientations form a natural group with identity, product and inverse.
+    This function thus has the following properties (here ``e`` is the identity,
+    ``*`` the product and ``-`` the inverse; ``a`` and ``b`` are arbitrary orientations):
+
+        reorient(pe_dir, e, e) == pe_dir
+        reorient(pe_dir, a, e) == reorient(pe_dir, a)
+        reorient(pe_dir, e, b) == reorient(pe_dir, -b)
+        reorient(pe_dir, a, b) == reorient(pe_dir, a * -b, e)
+
+    Therefore, this function accepts one or two orientations, and precomputed transforms
+    from a to b can be passed as the "source" orientation.
+
+    Passing only a source orientation will rotate into RAS:
+
+    >>> from nibabel.orientations import axcodes2ornt
+    >>> reorient_pedir("j", axcodes2ornt("RAS"))
+    'j'
+    >>> reorient_pedir("i", axcodes2ornt("PSL"))
+    'j-'
+
+    Passing only a target_ornt will rotate from RAS:
+
+    >>> reorient_pedir("j", source_ornt=None, target_ornt=axcodes2ornt("RAS"))
+    'j'
+    >>> reorient_pedir("i", source_ornt=None, target_ornt=axcodes2ornt("PSL"))
+    'k-'
+
+    Passing both will rotate from source to target.
+
+    >>> reorient_pedir("j", axcodes2ornt("LPS"), axcodes2ornt("AIR"))
+    'i-'
+
+    Passing a transform orientation as source_ornt will perform the transform,
+    and passing it as ``target_ornt`` will invert the transform:
+
+    >>> from nibabel.orientations import ornt_transform
+    >>> xfm = ornt_transform(axcodes2ornt("LPS"), axcodes2ornt("AIR"))
+    >>> reorient_pedir("j", xfm)
+    'i-'
+    >>> reorient_pedir("j", source_ornt=None, target_ornt=xfm)
+    'k-'
+    """
+    from nibabel.orientations import ornt_transform
+
+    if source_ornt is None:
+        source_ornt = [[0, 1], [1, 1], [2, 1]]
+    if target_ornt is None:
+        target_ornt = [[0, 1], [1, 1], [2, 1]]
+
+    xfm = ornt_transform(source_ornt, target_ornt).astype(int)  # shape: (3, 2)
+
+    directions = "ijk" if pe_dir[0] in "ijk" else "xyz"
+
+    source_axis = directions.index(pe_dir[0])
+    source_flip = pe_dir[1:] == "-"
+
+    axis_xfm = xfm[source_axis, :]  # shape: (2,)
+
+    target_axis = directions[axis_xfm[0]]
+    target_flip = source_flip ^ (axis_xfm[1] == -1)
+
+    return f"{target_axis}-" if target_flip else target_axis

sdcflows/workflows/fit/pepolar.py

@@ -92,7 +92,7 @@ def init_topup_wf(
 
     from ...utils.misc import front as _front
     from ...interfaces.epi import GetReadoutTime, SortPEBlips
-    from ...interfaces.utils import UniformGrid, PadSlices
+    from ...interfaces.utils import UniformGrid, PadSlices, ReorientImageAndMetadata
     from ...interfaces.bspline import TOPUPCoeffReorient
     from ..ancillary import init_brainextraction_wf
 
@@ -145,7 +145,7 @@ def init_topup_wf(
     setwise_avg = pe.Node(RobustAverage(num_threads=omp_nthreads), name="setwise_avg")
     # The core of the implementation
     # Feed the input images in LAS orientation, so FSL does not run funky reorientations
-    to_las = pe.Node(ReorientImage(target_orientation="LAS"), name="to_las")
+    to_las = pe.Node(ReorientImageAndMetadata(target_orientation="LAS"), name="to_las")
     topup = pe.Node(
         TOPUP(
             config=_pkg_fname("sdcflows", f"data/flirtsch/b02b0{'_quick' * sloppy}.cnf")
@@ -172,16 +172,19 @@ def init_topup_wf(
         (regrid, sort_pe_blips, [("out_data", "in_data")]),
         (readout_time, sort_pe_blips, [("readout_time", "readout_times"),
                                        ("pe_dir_fsl", "pe_dirs_fsl")]),
-        (sort_pe_blips, topup, [("readout_times", "readout_times"),
-                                ("pe_dirs_fsl", "encoding_direction")]),
-        (sort_pe_blips, fix_coeff, [(("pe_dirs", _front), "pe_dir")]),
+        (sort_pe_blips, topup, [("readout_times", "readout_times")]),
         (setwise_avg, fix_coeff, [("out_file", "fmap_ref")]),
         (sort_pe_blips, concat_blips, [("out_data", "in_files")]),
         (concat_blips, pad_blip_slices, [("out_file", "in_file")]),
         (pad_blip_slices, setwise_avg, [("out_file", "in_file")]),
         (setwise_avg, to_las, [("out_hmc_volumes", "in_file")]),
-        (to_las, topup, [("out_file", "in_file")]),
+        (sort_pe_blips, to_las, [("pe_dirs_fsl", "pe_dir")]),
+        (to_las, topup, [
+            ("out_file", "in_file"),
+            ("pe_dir", "encoding_direction"),
+        ]),
         (topup, fix_coeff, [("out_fieldcoef", "in_coeff")]),
+        (to_las, fix_coeff, [(("pe_dir", _front), "pe_dir")]),
         (topup, outputnode, [("out_jacs", "jacobians"),
                              ("out_mats", "xfms")]),
         (ref_average, brainextraction_wf, [("out_file", "inputnode.in_file")]),