fix: final revision of the base workflow, preparing for new API

oesteban · oesteban · commit 618be5b714e7 · 2020-11-01T15:13:44.000+01:00
diff --git a/sdcflows/interfaces/utils.py b/sdcflows/interfaces/utils.py
@@ -0,0 +1,99 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+"""
+Utilities.
+
+    .. testsetup::
+
+        >>> tmpdir = getfixture('tmpdir')
+        >>> tmp = tmpdir.chdir() # changing to a temporary directory
+        >>> nb.Nifti1Image(np.zeros((90, 90, 60)), None, None).to_filename(
+        ...     tmpdir.join('epi.nii.gz').strpath)
+
+"""
+
+from nipype import logging
+from nipype.interfaces.base import (
+    BaseInterfaceInputSpec,
+    TraitedSpec,
+    File,
+    traits,
+    SimpleInterface,
+    InputMultiObject,
+    OutputMultiObject,
+)
+
+LOGGER = logging.getLogger("nipype.interface")
+
+
+class _FlattenInputSpec(BaseInterfaceInputSpec):
+    in_data = InputMultiObject(
+        File(exists=True), mandatory=True, desc="list of input data",
+    )
+    in_meta = InputMultiObject(
+        traits.DictStrAny, mandatory=True, desc="list of metadata",
+    )
+    max_trs = traits.Int(50, usedefault=True, desc="only pick first TRs")
+
+
+class _FlattenOutputSpec(TraitedSpec):
+    out_list = OutputMultiObject(
+        traits.Tuple(File(exists=True), traits.DictStrAny,), desc="list of output files"
+    )
+    out_data = OutputMultiObject(File(exists=True))
+    out_meta = OutputMultiObject(traits.DictStrAny)
+
+
+class Flatten(SimpleInterface):
+    """Flatten a list of 3D and 4D files (and metadata)."""
+
+    input_spec = _FlattenInputSpec
+    output_spec = _FlattenOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results["out_list"] = _flatten(
+            zip(self.inputs.inlist, self.inputs.in_meta),
+            max_trs=self.inputs.max_trs,
+            out_dir=runtime.cwd,
+        )
+
+        # Unzip out_data, out_meta outputs.
+        self._results["out_data"], self._results["out_meta"] = zip(
+            *self._results["out_list"]
+        )
+        return runtime
+
+
+def _flatten(inlist, max_trs=50, out_dir=None):
+    """
+    Split the input EPIs and generate a flattened list with corresponding metadata.
+
+    Inputs
+    ------
+    inlist : :obj:`list` of :obj:`tuple`
+        List of pairs (filepath, metadata)
+    max_trs : int
+        Index of frame after which all volumes will be discarded
+        from the input EPI images.
+
+    """
+    from pathlib import Path
+    import nibabel as nb
+
+    out_dir = Path(out_dir) if out_dir is not None else Path()
+
+    output = []
+    for i, (path, meta) in enumerate(inlist):
+        img = nb.load(path)
+        if len(img.shape) == 3:
+            output.append((path, meta))
+        else:
+            splitnii = nb.four_to_three(img.slicer[:, :, :, :max_trs])
+            stem = Path(path).name.rpartition(".nii")[0]
+
+            for j, nii in enumerate(splitnii):
+                out_name = (out_dir / f"{stem}_idx-{j:03}.nii.gz").absolute()
+                nii.to_filename(out_name)
+                output.append((str(out_name), meta))
+
+    return output
diff --git a/sdcflows/models/fieldmap.py b/sdcflows/models/fieldmap.py
@@ -37,7 +37,7 @@
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
 
 
-def init_fmap_wf(omp_nthreads, mode="phase-diff", name="fmap_wf"):
+def init_fmap_wf(omp_nthreads, mode="phasediff", name="fmap_wf"):
     """
     Estimate the fieldmap based on a field-mapping MRI acquisition.
 
@@ -111,7 +111,7 @@ def init_fmap_wf(omp_nthreads, mode="phase-diff", name="fmap_wf"):
     ])
     # fmt: on
 
-    if mode == "phase-diff":
+    if mode == "phasediff":
         workflow.__desc__ = """\
 A *B<sub>0</sub>* nonuniformity map (or *fieldmap*) was estimated from the
 phase-drift map(s) measure with two consecutive GRE (gradient-recall echo)
diff --git a/sdcflows/models/pepolar.py b/sdcflows/models/pepolar.py
@@ -53,9 +53,9 @@ def init_topup_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
 
     Outputs
     -------
-    fieldmap : :obj:`str`
+    fmap : :obj:`str`
         The path of the estimated fieldmap.
-    corrected : :obj:`str`
+    fmap_ref : :obj:`str`
         The path of an unwarped conversion of files in ``in_data``.
 
     """
@@ -74,8 +74,8 @@ def init_topup_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
     outputnode = pe.Node(
         niu.IdentityInterface(
             fields=[
-                "corrected",
-                "fieldmap",
+                "fmap_ref",
+                "fmap",
                 "coefficients",
                 "jacobians",
                 "xfms",
@@ -87,7 +87,9 @@ def init_topup_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
 
     concat_blips = pe.Node(MergeSeries(), name="concat_blips")
     readout_time = pe.MapNode(
-        GetReadoutTime(), name="readout_time", iterfield=["metadata", "in_file"],
+        GetReadoutTime(),
+        name="readout_time",
+        iterfield=["metadata", "in_file"],
         run_without_submitting=True,
     )
 
@@ -106,8 +108,8 @@ def init_topup_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
         (inputnode, topup, [(("metadata", _pe2fsl), "encoding_direction")]),
         (readout_time, topup, [("readout_time", "readout_times")]),
         (concat_blips, topup, [("out_file", "in_file")]),
-        (topup, outputnode, [("out_corrected", "corrected"),
-                             ("out_field", "fieldmap"),
+        (topup, outputnode, [("out_corrected", "fmap_ref"),
+                             ("out_field", "fmap"),
                              ("out_fieldcoef", "coefficients"),
                              ("out_jacs", "jacobians"),
                              ("out_mats", "xfms"),
@@ -118,7 +120,7 @@ def init_topup_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
     return workflow
 
 
-def init_3dQwarp_wf(pe_dir, omp_nthreads=1, name="pepolar_estimate_wf"):
+def init_3dQwarp_wf(omp_nthreads=1, name="pepolar_estimate_wf"):
     """
     Create the PEPOLAR field estimation workflow based on AFNI's ``3dQwarp``.
 
@@ -132,12 +134,10 @@ def init_3dQwarp_wf(pe_dir, omp_nthreads=1, name="pepolar_estimate_wf"):
             :simple_form: yes
 
             from sdcflows.models.pepolar import init_3dQwarp_wf
-            wf = init_3dQwarp_wf(pe_dir="j")
+            wf = init_3dQwarp_wf()
 
     Parameters
     ----------
-    pe_dir : :obj:`str`
-        PE direction (BIDS compatible)
     name : :obj:`str`
         Name for this workflow
     omp_nthreads : :obj:`int`
@@ -150,34 +150,66 @@ def init_3dQwarp_wf(pe_dir, omp_nthreads=1, name="pepolar_estimate_wf"):
 
     Outputs
     -------
-    fieldmap : :obj:`str`
+    fmap : :obj:`str`
         The path of the estimated fieldmap.
-    corrected : :obj:`str`
+    fmap_ref : :obj:`str`
         The path of an unwarped conversion of the first element of ``in_data``.
 
     """
     from nipype.interfaces import afni
     from niworkflows.interfaces import CopyHeader
-    from niworkflows.interfaces.registration import ANTSApplyTransformsRPT
+    from niworkflows.interfaces.fixes import (
+        FixHeaderRegistration as Registration,
+        FixHeaderApplyTransforms as ApplyTransforms,
+    )
+    from niworkflows.interfaces.freesurfer import StructuralReference
+    from niworkflows.func.util import init_enhance_and_skullstrip_bold_wf
+    from ..interfaces.utils import Flatten
 
     workflow = Workflow(name=name)
     workflow.__desc__ = f"""{_PEPOLAR_DESC} \
 with `3dQwarp` @afni (AFNI {''.join(['%02d' % v for v in afni.Info().version() or []])}).
 """
 
-    inputnode = pe.Node(niu.IdentityInterface(fields=["in_data"]), name="inputnode")
+    inputnode = pe.Node(niu.IdentityInterface(fields=["in_data", "metadata"]),
+                        name="inputnode")
 
     outputnode = pe.Node(
-        niu.IdentityInterface(fields=["fieldmap", "corrected"]), name="outputnode"
+        niu.IdentityInterface(fields=["fmap", "fmap_ref"]), name="outputnode"
+    )
+
+    flatten = pe.Node(Flatten(), name="flatten")
+    sort_pe = pe.Node(niu.Function(
+        function=_sorted_pe, output_names=["sorted", "qwarp_args"]),
+        name="sort_pe", run_without_submitting=True)
+
+    merge_pes = pe.MapNode(StructuralReference(
+        auto_detect_sensitivity=True,
+        initial_timepoint=1,
+        fixed_timepoint=True,  # Align to first image
+        intensity_scaling=True,
+        # 7-DOF (rigid + intensity)
+        no_iteration=True,
+        subsample_threshold=200,
+        out_file='template.nii.gz'),
+        name='merge_pes',
+        iterfield=["in_files"],
+    )
+
+    pe0_wf = init_enhance_and_skullstrip_bold_wf(omp_nthreads=omp_nthreads, name="pe0_wf")
+    pe1_wf = init_enhance_and_skullstrip_bold_wf(omp_nthreads=omp_nthreads, name="pe1_wf")
+
+    align_pes = pe.Node(
+        Registration(
+            from_file=_pkg_fname("sdcflows", "data/translation_rigid.json"),
+            output_warped_image=True,
+        ),
+        name="align_pes",
+        n_procs=omp_nthreads,
     )
 
     qwarp = pe.Node(
         afni.QwarpPlusMinus(
-            args={
-                "i": "-noYdis -noZdis",
-                "j": "-noXdis -noZdis",
-                "k": "-noXdis -noYdis",
-            }[pe_dir[0]],
             blur=[-1, -1],
             environ={"OMP_NUM_THREADS": f"{omp_nthreads}"},
             minpatch=9,
@@ -194,9 +226,8 @@ def init_3dQwarp_wf(pe_dir, omp_nthreads=1, name="pepolar_estimate_wf"):
     cphdr_warp = pe.Node(CopyHeader(), name="cphdr_warp", mem_gb=0.01)
 
     unwarp_reference = pe.Node(
-        ANTSApplyTransformsRPT(
+        ApplyTransforms(
             dimension=3,
-            generate_report=False,
             float=True,
             interpolation="LanczosWindowedSinc",
         ),
@@ -205,16 +236,25 @@ def init_3dQwarp_wf(pe_dir, omp_nthreads=1, name="pepolar_estimate_wf"):
 
     # fmt: off
     workflow.connect([
-        (inputnode, qwarp, [(("in_data", _front), "in_file"),
-                            (("in_data", _last), "base_file")]),
+        (inputnode, flatten, [("in_data", "in_data"),
+                              ("metadata", "in_meta")]),
+        (flatten, sort_pe, [("out_list", "inlist")]),
+        (sort_pe, qwarp, [("qwarp_args", "args")]),
+        (sort_pe, merge_pes, [("sorted", "in_files")]),
+        (merge_pes, pe0_wf, [(("out_file", _front), "inputnode.in_file")]),
+        (merge_pes, pe1_wf, [(("out_file", _last), "inputnode.in_file")]),
+        (pe0_wf, align_pes, [("outputnode.skull_stripped_file", "fixed_image")]),
+        (pe1_wf, align_pes, [("outputnode.skull_stripped_file", "moving_image")]),
+        (pe0_wf, qwarp, [("outputnode.skull_stripped_file", "in_file")]),
+        (align_pes, qwarp, [("warped_image", "base_file")]),
         (inputnode, cphdr_warp, [(("in_data", _front), "hdr_file")]),
         (qwarp, cphdr_warp, [("source_warp", "in_file")]),
         (cphdr_warp, to_ants, [("out_file", "in_file")]),
         (to_ants, unwarp_reference, [("out", "transforms")]),
         (inputnode, unwarp_reference, [("in_reference", "reference_image"),
                                        ("in_reference", "input_image")]),
-        (unwarp_reference, outputnode, [("output_image", "corrected")]),
-        (to_ants, outputnode, [("out", "fieldmap")]),
+        (unwarp_reference, outputnode, [("output_image", "fmap_ref")]),
+        (to_ants, outputnode, [("out", "fmap")]),
     ])
     # fmt: on
     return workflow
@@ -238,11 +278,35 @@ def _fix_hdr(in_file, newpath=None):
 def _pe2fsl(metadata):
     """Convert ijk notation to xyz."""
     return [
-        m["PhaseEncodingDirection"].replace("i", "x").replace("j", "y").replace("k", "z")
+        m["PhaseEncodingDirection"]
+        .replace("i", "x")
+        .replace("j", "y")
+        .replace("k", "z")
         for m in metadata
     ]
 
 
+def _sorted_pe(inlist):
+    out_ref = [inlist[0][0]]
+    out_opp = []
+
+    ref_pe = inlist[0][1]["PhaseEncodingDirection"]
+    for d, m in inlist[1:]:
+        pe = m["PhaseEncodingDirection"]
+        if pe == ref_pe:
+            out_ref.append(d)
+        elif pe[0] == ref_pe[0]:
+            out_opp.append(d)
+        else:
+            raise ValueError("Cannot handle orthogonal PE encodings.")
+
+    return [out_ref, out_opp], {
+        "i": "-noYdis -noZdis",
+        "j": "-noXdis -noZdis",
+        "k": "-noXdis -noYdis",
+    }[ref_pe[0]]
+
+
 def _front(inlist):
     if isinstance(inlist, (list, tuple)):
         return inlist[0]
diff --git a/sdcflows/models/syn.py b/sdcflows/models/syn.py
@@ -81,15 +81,11 @@ def init_syn_sdc_wf(omp_nthreads, epi_pe=None, atlas_threshold=3, name="syn_sdc_
 
     Outputs
     -------
-    out_reference
-        the ``in_reference`` image after unwarping
-    out_reference_brain
-        the ``in_reference_brain`` image after unwarping
-    out_warp
+    fmap
         the corresponding :abbr:`DFM (displacements field map)` compatible with
         ANTs
-    out_mask
-        mask of the unwarped input file
+    fmap_ref
+        the ``in_reference`` image after unwarping
 
     References
     ----------
@@ -135,10 +131,7 @@ def init_syn_sdc_wf(omp_nthreads, epi_pe=None, atlas_threshold=3, name="syn_sdc_
         name="inputnode",
     )
     outputnode = pe.Node(
-        niu.IdentityInterface(
-            ["out_reference", "out_reference_brain", "out_mask", "out_warp"]
-        ),
-        name="outputnode",
+        niu.IdentityInterface(["fmap", "fmap_ref", "fmap_mask"]), name="outputnode",
     )
 
     # Collect predefined data
@@ -224,16 +217,13 @@ def init_syn_sdc_wf(omp_nthreads, epi_pe=None, atlas_threshold=3, name="syn_sdc_
         (inputnode, syn, [("in_reference_brain", "moving_image")]),
         (t1_2_ref, syn, [("output_image", "fixed_image")]),
         (fixed_image_masks, syn, [("out", "fixed_image_masks")]),
-        (syn, outputnode, [("forward_transforms", "out_warp")]),
+        (syn, outputnode, [("forward_transforms", "fmap")]),
         (syn, unwarp_ref, [("forward_transforms", "transforms")]),
         (inputnode, unwarp_ref, [("in_reference", "reference_image"),
                                  ("in_reference", "input_image")]),
         (unwarp_ref, skullstrip_bold_wf, [
             ("output_image", "inputnode.in_file")]),
-        (unwarp_ref, outputnode, [("output_image", "out_reference")]),
-        (skullstrip_bold_wf, outputnode, [
-            ("outputnode.skull_stripped_file", "out_reference_brain"),
-            ("outputnode.mask_file", "out_mask")]),
+        (unwarp_ref, outputnode, [("output_image", "fmap_ref")]),
     ])
     # fmt: on
 
diff --git a/sdcflows/workflows/base.py b/sdcflows/workflows/base.py
diff --git a/sdcflows/workflows/tests/test_base.py b/sdcflows/workflows/tests/test_base.py
