enh: general refactor

- Connections to the particular ``dwi_file`` are done within the
  particular run's preproc workflow, using ``KeySelect`` to demux
  the available fieldmap options.
  This relies on implementing nipreps/sdcflows#147 and
  nipreps/sdcflows#148.
- Minimizes the overhead in ``dmriprep/workflows/base.py``
- Run black

Author: oesteban

dmriprep/workflows/base.py

@@ -287,6 +287,7 @@ def init_single_subject_wf(subject_id):
         return workflow
 
     from .dwi.base import init_dwi_preproc_wf
+
     # Append the dMRI section to the existing anatomical excerpt
     # That way we do not need to stream down the number of DWI datasets
     anat_preproc_wf.__postdesc__ = (
@@ -300,77 +301,87 @@ def init_single_subject_wf(subject_id):
 and a *b=0* average for reference to the subsequent steps of preprocessing was calculated.
 """
     )
-    dwi_preproc_list = []
-    for dwi_file in subject_data["dwi"]:
-        dwi_preproc_wf = init_dwi_preproc_wf(dwi_file)
-        dwi_preproc_list.append(dwi_preproc_wf)
 
-    dwi_preproc_list_wf = pe.Node(niu.IdentityInterface(fields=["dwi_workflows"]),
-                                  name="dwi_preproc_list_wf")
-    dwi_preproc_list_wf.iterables = [("dwi_workflows", dwi_preproc_list)]
+    # SDC Step 0: Determine whether fieldmaps can/should be estimated
+    fmap_estimators = None
+    if "fieldmap" not in config.workflow.ignore:
+        from sdcflows import fieldmaps as fm
+        from sdcflows.utils.wrangler import find_estimators
+        from sdcflows.workflows.base import init_fmap_preproc_wf
 
-    # fmt: off
-    workflow.connect([
-        (anat_preproc_wf, dwi_preproc_list_wf, [
-            ("outputnode.t1w_preproc", "inputnode.t1w_preproc"),
-            ("outputnode.t1w_mask", "inputnode.t1w_mask"),
-            ("outputnode.t1w_dseg", "inputnode.t1w_dseg"),
-            ("outputnode.t1w_aseg", "inputnode.t1w_aseg"),
-            ("outputnode.t1w_aparc", "inputnode.t1w_aparc"),
-            ("outputnode.t1w_tpms", "inputnode.t1w_tpms"),
-            ("outputnode.template", "inputnode.template"),
-            ("outputnode.anat2std_xfm", "inputnode.anat2std_xfm"),
-            ("outputnode.std2anat_xfm", "inputnode.std2anat_xfm"),
-            # Undefined if --fs-no-reconall, but this is safe
-            ("outputnode.subjects_dir", "inputnode.subjects_dir"),
-            ("outputnode.t1w2fsnative_xfm", "inputnode.t1w2fsnative_xfm"),
-            ("outputnode.fsnative2t1w_xfm", "inputnode.fsnative2t1w_xfm"),
-        ]),
-        (bids_info, dwi_preproc_list_wf, [("subject", "inputnode.subject_id")]),
-    ])
-    # fmt: on
+        # SDC Step 1: Run basic heuristics to identify available data for fieldmap estimation
+        fmap_estimators = find_estimators(config.execution.layout)
 
-    if "fieldmap" in config.workflow.ignore:
-        return workflow
+        # Add fieldmap-less estimators
+        if not fmap_estimators and config.workflow.use_syn:
+            # estimators = [fm.FieldmapEstimation()]
+            raise NotImplementedError
 
-    from sdcflows import fieldmaps as fm
-    from sdcflows.utils.wrangler import find_estimators
-    from sdcflows.workflows.base import init_fmap_preproc_wf
+    # Nuts and bolts: initialize individual run's pipeline
+    dwi_preproc_list = []
+    for dwi_file in subject_data["dwi"]:
+        dwi_preproc_wf = init_dwi_preproc_wf(
+            dwi_file,
+            has_fieldmap=bool(fmap_estimators),
+        )
 
-    # SDCFlows connection
-    # Step 1: Run basic heuristics to identify available data for fieldmap estimation
-    estimators = find_estimators(config.execution.layout)
+        # fmt: off
+        workflow.connect([
+            (anat_preproc_wf, dwi_preproc_wf, [
+                ("outputnode.t1w_preproc", "inputnode.t1w_preproc"),
+                ("outputnode.t1w_mask", "inputnode.t1w_mask"),
+                ("outputnode.t1w_dseg", "inputnode.t1w_dseg"),
+                ("outputnode.t1w_aseg", "inputnode.t1w_aseg"),
+                ("outputnode.t1w_aparc", "inputnode.t1w_aparc"),
+                ("outputnode.t1w_tpms", "inputnode.t1w_tpms"),
+                ("outputnode.template", "inputnode.template"),
+                ("outputnode.anat2std_xfm", "inputnode.anat2std_xfm"),
+                ("outputnode.std2anat_xfm", "inputnode.std2anat_xfm"),
+                # Undefined if --fs-no-reconall, but this is safe
+                ("outputnode.subjects_dir", "inputnode.subjects_dir"),
+                ("outputnode.t1w2fsnative_xfm", "inputnode.t1w2fsnative_xfm"),
+                ("outputnode.fsnative2t1w_xfm", "inputnode.fsnative2t1w_xfm"),
+            ]),
+            (bids_info, dwi_preproc_wf, [("subject", "inputnode.subject_id")]),
+        ])
+        # fmt: on
+
+        # Keep a handle to each workflow
+        dwi_preproc_list.append(dwi_preproc_wf)
 
-    if not estimators and config.workflow.use_syn:  # Add fieldmap-less estimators
-        # estimators = [fm.FieldmapEstimation()]
-        raise NotImplementedError
+    if not fmap_estimators:
+        return workflow
 
-    # Step 2: Manually add further estimators (e.g., fieldmap-less)
+    # SDC Step 2: Manually add further estimators (e.g., fieldmap-less)
     fmap_wf = init_fmap_preproc_wf(
         debug=config.execution.debug,
-        estimators=estimators,
+        estimators=fmap_estimators,
         omp_nthreads=config.nipype.omp_nthreads,
         output_dir=str(output_dir),
         subject=subject_id,
     )
-    # fmt: off
-    workflow.connect([
-        (fmap_wf, dwi_preproc_list_wf, [
-            ("outputnode.fmap", "inputnode.fmap"),
-            ("outputnode.fmap_ref", "inputnode.fmap_ref"),
-            ("outputnode.fmap_coeff", "inputnode.fmap_coeff"),
-            ("outputnode.fmap_mask", "inputnode.fmap_mask"),
-        ]),
-    ])
-    # fmt: on
+
+    # TODO: Requires nipreps/sdcflows#147
+    for dwi_preproc_wf in dwi_preproc_list:
+        # fmt: off
+        workflow.connect([
+            (fmap_wf, dwi_preproc_wf, [
+                ("outputnode.fmap", "inputnode.fmap"),
+                ("outputnode.fmap_ref", "inputnode.fmap_ref"),
+                ("outputnode.fmap_coeff", "inputnode.fmap_coeff"),
+                ("outputnode.fmap_mask", "inputnode.fmap_mask"),
+                ("outputnode.fmap_id", "inputnode.fmap_id"),
+            ]),
+        ])
+        # fmt: on
+
     # Overwrite ``out_path_base`` of sdcflows's DataSinks
     for node in fmap_wf.list_node_names():
         if node.split(".")[-1].startswith("ds_"):
             fmap_wf.get_node(node).interface.out_path_base = "dmriprep"
-    workflow.add_nodes([fmap_wf])
 
     # Step 3: Manually connect PEPOLAR
-    for estimator in estimators:
+    for estimator in fmap_estimators:
         if estimator.method != fm.EstimatorType.PEPOLAR:
             continue
 
@@ -387,18 +398,18 @@ def init_single_subject_wf(subject_id):
             raise NotImplementedError
             # from niworkflows.interfaces.utility import KeySelect
             # est_id = estimator.bids_id
-            # fmap_select = pe.MapNode(
+            # estim_select = pe.MapNode(
             #     KeySelect(fields=["metadata", "dwi_reference", "dwi_mask", "gradients_rasb",]),
             #     name=f"fmap_select_{est_id}",
             #     run_without_submitting=True,
             #     iterfields=["key"]
             # )
-            # fmap_select.inputs.key = [
+            # estim_select.inputs.key = [
             #     str(s.path) for s in estimator.sources if s.suffix in ("epi", "dwi", "sbref")
             # ]
             # # fmt:off
             # workflow.connect([
-            #     (referencenode, fmap_select, [("dwi_file", "keys"),
+            #     (referencenode, estim_select, [("dwi_file", "keys"),
             #                                   ("metadata", "metadata"),
             #                                   ("dwi_reference", "dwi_reference"),
             #                                   ("gradients_rasb", "gradients_rasb")]),

dmriprep/workflows/dwi/base.py

@@ -83,7 +83,7 @@ def init_dwi_preproc_wf(dwi_file, has_fieldmap=False):
                 "dwi_file",
                 "in_bvec",
                 "in_bval",
-                # From fmap
+                # From SDCFlows
                 "fmap",
                 "fmap_ref",
                 "fmap_coeff",
@@ -129,10 +129,6 @@ def init_dwi_preproc_wf(dwi_file, has_fieldmap=False):
                                      ("in_bval", "in_bval")]),
         (inputnode, dwi_reference_wf, [("dwi_file", "inputnode.dwi_file")]),
         (gradient_table, dwi_reference_wf, [("b0_ixs", "inputnode.b0_ixs")]),
-        #outputnode, [
-        #    ("outputnode.ref_image", "dwi_reference"),
-        #    ("outputnode.dwi_mask", "dwi_mask"),
-        #]),
         (gradient_table, outputnode, [("out_rasb", "gradients_rasb")]),
     ])
     # fmt: on
@@ -184,35 +180,68 @@ def _bold_reg_suffix(fallback):
         ])
         # fmt: on
 
-    if has_fieldmap:
-        from sdcflows.workflows.apply.registration import init_coeff2epi_wf
-        from sdcflows.workflows.apply.correction import init_unwarp_wf
-
-        coeff2epi_wf = init_coeff2epi_wf(
-            omp_nthreads=config.nipype.omp_nthreads, write_coeff=True
-        )
-        unwarp_wf = init_unwarp_wf(omp_nthreads=config.nipype.omp_nthreads)
-
-        # fmt: off
-        workflow.connect([
-            (dwi_reference_wf, coeff2epi_wf, [
-                ("outputnode.ref_image", "inputnode.target_ref"),
-                ("outputnode.dwi_mask", "inputnode.target_mask")]),
-            (coeff2epi_wf, unwarp_wf, [
-                ("outputnode.fmap_coeff", "inputnode.fmap_coeff")])
-        ])
-        # fmt: on
-
     # REPORTING ############################################################
     reportlets_wf = init_reportlets_wf(str(config.execution.output_dir))
     # fmt: off
     workflow.connect([
         (inputnode, reportlets_wf, [("dwi_file", "inputnode.source_file")]),
         (dwi_reference_wf, reportlets_wf, [
-            ("outputnode.ref_image", "inputnode.dwi_ref"),
-            ("outputnode.dwi_mask", "inputnode.dwi_mask"),
             ("outputnode.validation_report", "inputnode.validation_report"),
         ]),
+        (outputnode, reportlets_wf, [
+            ("dwi_reference", "inputnode.dwi_ref"),
+            ("dwi_mask", "inputnode.dwi_mask"),
+        ]),
+    ])
+    # fmt: on
+
+    if not has_fieldmap:
+        # fmt: off
+        workflow.connect([
+            (dwi_reference_wf, outputnode, [("outputnode.ref_image", "dwi_reference"),
+                                            ("outputnode.dwi_mask", "dwi_mask")]),
+        ])
+        # fmt: on
+        return workflow
+
+    from niworkflows.interfaces.utility import KeySelect
+    from sdcflows.workflows.apply.registration import init_coeff2epi_wf
+    from sdcflows.workflows.apply.correction import init_unwarp_wf
+
+    # TODO: Requires nipreps/sdcflows#148
+    # from sdcflows.utils.fieldmap import get_identifier
+
+    coeff2epi_wf = init_coeff2epi_wf(
+        omp_nthreads=config.nipype.omp_nthreads, write_coeff=True
+    )
+    unwarp_wf = init_unwarp_wf(omp_nthreads=config.nipype.omp_nthreads)
+    unwarp_wf.inputs.inputnode.metadata = layout.get_metadata(dwi_file)
+
+    output_select = pe.Node(
+        KeySelect(fields=["fmap", "fmap_ref", "fmap_coeff", "fmap_mask"]),
+        name="output_select",
+        run_without_submitting=True,
+    )
+    # output_select.inputs.key = get_identifier(dwi_file)
+
+    # fmt: off
+    workflow.connect([
+        (inputnode, output_select, [("fmap", "fmap"),
+                                    ("fmap_ref", "fmap_ref"),
+                                    ("fmap_coeff", "fmap_coeff"),
+                                    ("fmap_mask", "fmap_mask"),
+                                    ("fmap_id", "keys")]),
+        (output_select, coeff2epi_wf, [
+            ("fmap_ref", "inputnode.fmap_ref"),
+            ("fmap_coeff", "inputnode.fmap_coeff"),
+            ("fmap_mask", "inputnode.fmap_mask")]),
+        (dwi_reference_wf, coeff2epi_wf, [
+            ("outputnode.ref_image", "inputnode.target_ref"),
+            ("outputnode.dwi_mask", "inputnode.target_mask")]),
+        (dwi_reference_wf, unwarp_wf, [("outputnode.ref_image", "distorted")]),
+        (coeff2epi_wf, unwarp_wf, [
+            ("outputnode.fmap_coeff", "inputnode.fmap_coeff")]),
+        (unwarp_wf, outputnode, [("outputnode.corrected", "dwi_reference")]),
     ])
     # fmt: on