ENH: Refactor the distortion estimation workflow

Following nipreps#40, rolls back the changes most involved with PyBIDS also
making the API evolve less abruptly.

The PR builds on top of nipreps#52, and closes nipreps#16.

This PR does not address #20, nor nipreps#19 and nipreps#21 although it is setting the
necessary stepping stones.

Author: oesteban

sdcflows/workflows/base.py

@@ -1,16 +1,12 @@
 # emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """SDC workflows coordination."""
-from collections import defaultdict
-
 from nipype.pipeline import engine as pe
 from nipype.interfaces import utility as niu
 from nipype import logging
 
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
 
-# Fieldmap workflows
-from .pepolar import init_pepolar_unwarp_wf
 
 LOGGER = logging.getLogger('nipype.workflow')
 FMAP_PRIORITY = {
@@ -23,12 +19,12 @@
 DEFAULT_MEMORY_MIN_GB = 0.01
 
 
-def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
+def init_sdc_estimate_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
     """
     Build a :abbr:`SDC (susceptibility distortion correction)` workflow.
 
-    This workflow implements the heuristics to choose a
-    :abbr:`SDC (susceptibility distortion correction)` strategy.
+    This workflow implements the heuristics to choose an estimation
+    methodology for :abbr:`SDC (susceptibility distortion correction)`.
     When no field map information is present within the BIDS inputs,
     the EXPERIMENTAL "fieldmap-less SyN" can be performed, using
     the ``--use-syn`` argument. When ``--force-syn`` is specified,
@@ -41,10 +37,14 @@ def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
 
     Parameters
     ----------
-    epi_file : pybids.BIDSFile
-        A BIDSFile object representing an distorted image acquired
-        with an :abbr:`EPI (echo-planar imaging)` MR scheme
-        (i.e., suffix ``bold``, ``sbref`` or ``dwi``).
+    fmaps : list of pybids dicts
+        A list of dictionaries with the available fieldmaps
+        (and their metadata using the key ``'metadata'`` for the
+        case of :abbr:`PEPOLAR (Phase-Encoding POLARity)` fieldmaps).
+    epi_meta : dict
+        BIDS metadata dictionary corresponding to the
+        :abbr:`EPI (echo-planar imaging)` run (i.e., suffix ``bold``,
+        ``sbref``, or ``dwi``) for which the fieldmap is being estimated.
     omp_nthreads : int
         Maximum number of threads an individual process may use
     debug : bool
@@ -54,30 +54,23 @@ def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
     ------
     epi_file
         A reference image calculated at a previous stage
-    ref_brain
+    epi_brain
         Same as above, but brain-masked
-    ref_mask
+    epi_mask
         Brain mask for the run
-    t1_brain
+    t1w_brain
         T1w image, brain-masked, for the fieldmap-less SyN method
     std2anat_xfm
-        List of standard-to-T1w transforms generated during spatial
+        Standard-to-T1w transform generated during spatial
         normalization (only for the fieldmap-less SyN method).
-    template : str
-        Name of template from which prior knowledge will be mapped
-        into the subject's T1w reference
-        (only for the fieldmap-less SyN method)
-    templates : str
-        Name of templates that index the ``std2anat_xfm`` input list
-        (only for the fieldmap-less SyN method).
 
     Outputs
     -------
     epi_file
         An unwarped BOLD reference
-    bold_mask
+    epi_mask
         The corresponding new mask after unwarping
-    bold_ref_brain
+    epi_brain
         Brain-extracted, unwarped BOLD reference
     out_warp
         The deformation field to unwarp the susceptibility distortions
@@ -95,14 +88,13 @@ def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
     # TODO: To be removed (filter out unsupported fieldmaps):
     fmaps = [fmap for fmap in fmaps if fmap['suffix'] in FMAP_PRIORITY]
 
-    workflow = Workflow(name='sdc_wf' if fmaps else 'sdc_bypass_wf')
+    workflow = Workflow(name='sdc_estimate_wf' if fmaps else 'sdc_bypass_wf')
     inputnode = pe.Node(niu.IdentityInterface(
-        fields=['epi_file', 'ref_brain', 'ref_mask',
-                't1_brain', 'std2anat_xfm', 'template', 'templates']),
+        fields=['epi_file', 'epi_brain', 'epi_mask', 't1w_brain', 'std2anat_xfm']),
         name='inputnode')
 
     outputnode = pe.Node(niu.IdentityInterface(
-        fields=['output_ref', 'ref_mask', 'ref_brain',
+        fields=['output_ref', 'epi_mask', 'epi_brain',
                 'out_warp', 'syn_ref', 'method']),
         name='outputnode')
 
@@ -116,8 +108,8 @@ def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
         outputnode.inputs.method = 'None'
         workflow.connect([
             (inputnode, outputnode, [('epi_file', 'output_ref'),
-                                     ('ref_mask', 'ref_mask'),
-                                     ('ref_brain', 'ref_brain')]),
+                                     ('epi_mask', 'epi_mask'),
+                                     ('epi_brain', 'epi_brain')]),
         ])
         return workflow
 
@@ -127,109 +119,113 @@ def init_sdc_wf(fmaps, epi_meta, omp_nthreads=1, debug=False, ignore=None):
 accurate co-registration with the anatomical reference.
 """
 
+    # In case there are multiple fieldmaps prefer EPI
+    fmaps.sort(key=lambda fmap: FMAP_PRIORITY[fmap['suffix']])
+    fmap = fmaps[0]
+
     # PEPOLAR path
     if 'epi' in fmaps:
+        from .pepolar import init_pepolar_unwarp_wf, check_pes
         outputnode.inputs.method = 'PEB/PEPOLAR (phase-encoding based / PE-POLARity)'
+
+        # Filter out EPI fieldmaps to be used
+        fmaps_epi = [(epi.path, epi.get_metadata()['PhaseEncodingDirection'])
+                     for epi in fmaps['epi']]
+
+        # Find matched PE directions
+        matched_pe = check_pes(fmaps_epi, epi_meta['PhaseEncodingDirection'])
+
         # Get EPI polarities and their metadata
         sdc_unwarp_wf = init_pepolar_unwarp_wf(
-            bold_meta=epi_file.get_metadata(),
-            epi_fmaps=[(fmap, fmap.get_metadata()["PhaseEncodingDirection"])
-                       for fmap in fmaps['epi']],
-            omp_nthreads=omp_nthreads,
-            name='pepolar_unwarp_wf')
+            matched_pe=matched_pe,
+            omp_nthreads=omp_nthreads)
+        sdc_unwarp_wf.inputs.inputnode.epi_pe_dir = epi_meta['PhaseEncodingDirection']
+        sdc_unwarp_wf.inputs.inputnode.fmaps_epi = fmaps_epi
 
         workflow.connect([
             (inputnode, sdc_unwarp_wf, [
                 ('epi_file', 'inputnode.in_reference'),
-                ('bold_mask', 'inputnode.in_mask'),
-                ('bold_ref_brain', 'inputnode.in_reference_brain')]),
+                ('epi_brain', 'inputnode.in_reference_brain'),
+                ('epi_mask', 'inputnode.in_mask')]),
         ])
 
     # FIELDMAP path
-    # elif 'fieldmap' in fmaps:
-    #     # Import specific workflows here, so we don't break everything with one
-    #     # unused workflow.
-    #     suffices = {f.suffix for f in fmaps['fieldmap']}
-    #     if 'fieldmap' in suffices:
-    #         from .fmap import init_fmap_wf
-    #         outputnode.inputs.method = 'FMB (fieldmap-based)'
-    #         fmap_estimator_wf = init_fmap_wf(
-    #             omp_nthreads=omp_nthreads,
-    #             fmap_bspline=False)
-    #         # set inputs
-    #         fmap_estimator_wf.inputs.inputnode.fieldmap = fmap['fieldmap']
-    #         fmap_estimator_wf.inputs.inputnode.magnitude = fmap['magnitude']
-
-    #     if fmap['suffix'] == 'phasediff':
-    #         from .phdiff import init_phdiff_wf
-    #         fmap_estimator_wf = init_phdiff_wf(omp_nthreads=omp_nthreads)
-    #         # set inputs
-    #         fmap_estimator_wf.inputs.inputnode.phasediff = fmap['phasediff']
-    #         fmap_estimator_wf.inputs.inputnode.magnitude = [
-    #             fmap_ for key, fmap_ in sorted(fmap.items())
-    #             if key.startswith("magnitude")
-    #         ]
-
-    #     sdc_unwarp_wf = init_sdc_unwarp_wf(
-    #         omp_nthreads=omp_nthreads,
-    #         fmap_demean=fmap_demean,
-    #         debug=debug,
-    #         name='sdc_unwarp_wf')
-    #     sdc_unwarp_wf.inputs.inputnode.metadata = bold_meta
-
-    #     workflow.connect([
-    #         (inputnode, sdc_unwarp_wf, [
-    #             ('epi_file', 'inputnode.in_reference'),
-    #             ('bold_ref_brain', 'inputnode.in_reference_brain'),
-    #             ('bold_mask', 'inputnode.in_mask')]),
-    #         (fmap_estimator_wf, sdc_unwarp_wf, [
-    #             ('outputnode.fmap', 'inputnode.fmap'),
-    #             ('outputnode.fmap_ref', 'inputnode.fmap_ref'),
-    #             ('outputnode.fmap_mask', 'inputnode.fmap_mask')]),
-    #     ])
-
-    # # FIELDMAP-less path
-    # if any(fm['suffix'] == 'syn' for fm in fmaps):
-    #     # Select template
-    #     sdc_select_std = pe.Node(KeySelect(
-    #         fields=['std2anat_xfm']),
-    #         name='sdc_select_std', run_without_submitting=True)
-
-    #     syn_sdc_wf = init_syn_sdc_wf(
-    #         bold_pe=bold_meta.get('PhaseEncodingDirection', None),
-    #         omp_nthreads=omp_nthreads)
-
-    #     workflow.connect([
-    #         (inputnode, sdc_select_std, [
-    #             ('template', 'key'),
-    #             ('templates', 'keys'),
-    #             ('std2anat_xfm', 'std2anat_xfm')]),
-    #         (sdc_select_std, syn_sdc_wf, [
-    #             ('std2anat_xfm', 'inputnode.std2anat_xfm')]),
-    #         (inputnode, syn_sdc_wf, [
-    #             ('t1_brain', 'inputnode.t1_brain'),
-    #             ('epi_file', 'inputnode.epi_file'),
-    #             ('bold_ref_brain', 'inputnode.bold_ref_brain'),
-    #             ('template', 'inputnode.template')]),
-    #     ])
-
-    #     # XXX Eliminate branch when forcing isn't an option
-    #     if fmap['suffix'] == 'syn':  # No fieldmaps, but --use-syn
-    #         outputnode.inputs.method = 'FLB ("fieldmap-less", SyN-based)'
-    #         sdc_unwarp_wf = syn_sdc_wf
-    #     else:  # --force-syn was called when other fieldmap was present
-    #         sdc_unwarp_wf.__desc__ = None
-    #         workflow.connect([
-    #             (syn_sdc_wf, outputnode, [
-    #                 ('outputnode.out_reference', 'syn_bold_ref')]),
-    #         ])
+    elif 'fieldmap' in fmaps:
+        from .unwarp import init_sdc_unwarp_wf
+        # Import specific workflows here, so we don't break everything with one
+        # unused workflow.
+        suffices = {f.suffix for f in fmaps['fieldmap']}
+        if 'fieldmap' in suffices:
+            from .fmap import init_fmap_wf
+            outputnode.inputs.method = 'FMB (fieldmap-based)'
+            fmap_wf = init_fmap_wf(
+                omp_nthreads=omp_nthreads,
+                fmap_bspline=False)
+            # set inputs
+            fmap_wf.inputs.inputnode.magnitude = fmap['magnitude']
+            fmap_wf.inputs.inputnode.fieldmap = fmap['fieldmap']
+        elif 'phasediff' in suffices:
+            from .phdiff import init_phdiff_wf
+            fmap_wf = init_phdiff_wf(omp_nthreads=omp_nthreads)
+            # set inputs
+            fmap_wf.inputs.inputnode.phasediff = fmap['phasediff']
+            fmap_wf.inputs.inputnode.magnitude = [
+                fmap_ for key, fmap_ in sorted(fmap.items())
+                if key.startswith("magnitude")
+            ]
+        else:
+            raise ValueError('Fieldmaps of types %s are not supported' %
+                             ', '.join(['"%s"' % f for f in suffices]))
+
+        sdc_unwarp_wf = init_sdc_unwarp_wf(
+            omp_nthreads=omp_nthreads,
+            debug=debug,
+            name='sdc_unwarp_wf')
+        sdc_unwarp_wf.inputs.inputnode.metadata = epi_meta
+
+        workflow.connect([
+            (inputnode, sdc_unwarp_wf, [
+                ('epi_file', 'inputnode.in_reference'),
+                ('epi_brain', 'inputnode.in_reference_brain'),
+                ('epi_mask', 'inputnode.in_mask')]),
+            (fmap_wf, sdc_unwarp_wf, [
+                ('outputnode.fmap', 'inputnode.fmap'),
+                ('outputnode.fmap_ref', 'inputnode.fmap_ref'),
+                ('outputnode.fmap_mask', 'inputnode.fmap_mask')]),
+        ])
+
+    # FIELDMAP-less path
+    if any(fm['suffix'] == 'syn' for fm in fmaps):
+        from .syn import init_syn_sdc_wf
+        syn_sdc_wf = init_syn_sdc_wf(
+            bold_pe=epi_meta.get('PhaseEncodingDirection', None),
+            omp_nthreads=omp_nthreads)
+
+        workflow.connect([
+            (inputnode, syn_sdc_wf, [
+                ('t1w_brain', 'inputnode.t1w_brain'),
+                ('epi_file', 'inputnode.epi_file'),
+                ('epi_brain', 'inputnode.epi_brain'),
+                ('std2anat_xfm', 'inputnode.std2anat_xfm')]),
+        ])
+
+        # XXX Eliminate branch when forcing isn't an option
+        if fmap['suffix'] == 'syn':  # No fieldmaps, but --use-syn
+            outputnode.inputs.method = 'FLB ("fieldmap-less", SyN-based)'
+            sdc_unwarp_wf = syn_sdc_wf
+        else:  # --force-syn was called when other fieldmap was present
+            sdc_unwarp_wf.__desc__ = None
+            workflow.connect([
+                (syn_sdc_wf, outputnode, [
+                    ('outputnode.out_reference', 'syn_bold_ref')]),
+            ])
 
     workflow.connect([
         (sdc_unwarp_wf, outputnode, [
             ('outputnode.out_warp', 'out_warp'),
             ('outputnode.out_reference', 'epi_file'),
-            ('outputnode.out_reference_brain', 'bold_ref_brain'),
-            ('outputnode.out_mask', 'bold_mask')]),
+            ('outputnode.out_reference_brain', 'epi_brain'),
+            ('outputnode.out_mask', 'epi_mask')]),
     ])
 
     return workflow

sdcflows/workflows/syn.py

@@ -70,23 +70,23 @@ def init_syn_sdc_wf(omp_nthreads, bold_pe=None,
 
     Inputs
     ------
-    bold_ref
+    in_reference
         reference image
-    bold_ref_brain
+    in_reference_brain
         skull-stripped reference image
     template : str
         Name of template targeted by ``template`` output space
-    t1_brain
+    t1w_brain
         skull-stripped, bias-corrected structural image
     std2anat_xfm
         inverse registration transform of T1w image to MNI template
 
     Outputs
     -------
     out_reference
-        the ``bold_ref`` image after unwarping
+        the ``in_reference`` image after unwarping
     out_reference_brain
-        the ``bold_ref_brain`` image after unwarping
+        the ``in_reference_brain`` image after unwarping
     out_warp
         the corresponding :abbr:`DFM (displacements field map)` compatible with
         ANTs
@@ -127,8 +127,8 @@ def init_syn_sdc_wf(omp_nthreads, bold_pe=None,
 template [@fieldmapless3].
 """.format(ants_ver=Registration().version or '<ver>')
     inputnode = pe.Node(
-        niu.IdentityInterface(['bold_ref', 'bold_ref_brain', 'template',
-                               't1_brain', 'std2anat_xfm']),
+        niu.IdentityInterface(['in_reference', 'in_reference_brain', 'template',
+                               't1w_brain', 'std2anat_xfm']),
         name='inputnode')
     outputnode = pe.Node(
         niu.IdentityInterface(['out_reference', 'out_reference_brain',
@@ -184,28 +184,28 @@ def init_syn_sdc_wf(omp_nthreads, bold_pe=None,
     skullstrip_bold_wf = init_skullstrip_bold_wf()
 
     workflow.connect([
-        (inputnode, invert_t1w, [('t1_brain', 'in_file'),
-                                 ('bold_ref', 'ref_file')]),
-        (inputnode, ref_2_t1, [('bold_ref_brain', 'moving_image')]),
+        (inputnode, invert_t1w, [('t1w_brain', 'in_file'),
+                                 ('in_reference', 'ref_file')]),
+        (inputnode, ref_2_t1, [('in_reference_brain', 'moving_image')]),
         (invert_t1w, ref_2_t1, [('out_file', 'fixed_image')]),
-        (inputnode, t1_2_ref, [('bold_ref', 'reference_image')]),
+        (inputnode, t1_2_ref, [('in_reference', 'reference_image')]),
         (invert_t1w, t1_2_ref, [('out_file', 'input_image')]),
         (ref_2_t1, t1_2_ref, [('forward_transforms', 'transforms')]),
         (ref_2_t1, transform_list, [('forward_transforms', 'in1')]),
         (inputnode, transform_list, [
             ('std2anat_xfm', 'in2'),
             (('template', _prior_path), 'in3')]),
-        (inputnode, atlas_2_ref, [('bold_ref', 'reference_image')]),
+        (inputnode, atlas_2_ref, [('in_reference', 'reference_image')]),
         (transform_list, atlas_2_ref, [('out', 'transforms')]),
         (atlas_2_ref, threshold_atlas, [('output_image', 'in_file')]),
         (threshold_atlas, fixed_image_masks, [('out_file', 'in2')]),
-        (inputnode, syn, [('bold_ref_brain', 'moving_image')]),
+        (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, unwarp_ref, [('forward_transforms', 'transforms')]),
-        (inputnode, unwarp_ref, [('bold_ref', 'reference_image'),
-                                 ('bold_ref', 'input_image')]),
+        (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')]),