Simplify base workflow.

Author: tsalo

src/smriprep/workflows/anatomical.py

@@ -69,6 +69,7 @@
     init_ds_fs_segs_wf,
     init_ds_grayord_metrics_wf,
     init_ds_mask_wf,
+    init_ds_surface_masks_wf,
     init_ds_surface_metrics_wf,
     init_ds_surfaces_wf,
     init_ds_template_registration_wf,
@@ -78,7 +79,7 @@
 )
 from .surfaces import (
     init_anat_ribbon_wf,
-    init_cortex_mask_wf,
+    init_cortex_masks_wf,
     init_fsLR_reg_wf,
     init_gifti_morphometrics_wf,
     init_gifti_surfaces_wf,
@@ -1350,57 +1351,27 @@ def init_anat_fit_wf(
     if len(precomputed.get('cortex_mask', [])) < 2:
         LOGGER.info('ANAT Stage 11: Creating cortical surface mask')
 
-        # Merge outputs into a single list
-        merge_cortex_masks = pe.Node(
-            niu.Merge(2),
-            name='merge_cortex_masks',
-        )
-        workflow.connect([(merge_cortex_masks, outputnode, [('out', 'cortex_mask')])])
-
-        for i_hemi, hemi in enumerate(['L', 'R']):
-            select_midthickness = pe.Node(
-                niu.Select(index=i_hemi),
-                name=f'select_midthickness_{hemi}',
-            )
-            select_thickness = pe.Node(
-                niu.Select(index=i_hemi),
-                name=f'select_thickness_{hemi}',
-            )
-            cortex_mask_wf = init_cortex_mask_wf(name=f'cortex_mask_wf_{hemi}')
-            cortex_mask_wf.inputs.inputnode.hemi = hemi
-
-            workflow.connect([
-                (surfaces_buffer, select_midthickness, [('midthickness', 'inlist')]),
-                (surfaces_buffer, select_thickness, [('thickness', 'inlist')]),
-                (select_midthickness, cortex_mask_wf, [('out', 'inputnode.midthickness')]),
-                (select_thickness, cortex_mask_wf, [('out', 'inputnode.thickness')]),
-            ])  # fmt:skip
-
-            # Combine the inputs into a list
-            combine_inputs = pe.Node(
-                niu.Merge(2),
-                name=f'combine_inputs_{hemi}',
-            )
-            workflow.connect([
-                (select_midthickness, combine_inputs, [('out', 'in1')]),
-                (select_thickness, combine_inputs, [('out', 'in2')]),
-            ])  # fmt:skip
-
-            ds_cortex_mask_wf = init_ds_mask_wf(
-                bids_root=bids_root,
-                output_dir=output_dir,
-                mask_type='cortex',
-                name=f'ds_cortex_mask_wf_{hemi}',
-                extra_entities={'extension': '.label.gii', 'hemi': hemi},
-            )
+        cortex_masks_wf = init_cortex_masks_wf()
+        workflow.connect([
+            (surfaces_buffer, cortex_masks_wf, [
+                ('midthickness', 'inputnode.midthickness'),
+                ('thickness', 'inputnode.thickness'),
+            ]),
+        ])  # fmt:skip
 
-            workflow.connect([
-                (cortex_mask_wf, ds_cortex_mask_wf, [('outputnode.roi', 'inputnode.mask_file')]),
-                (combine_inputs, ds_cortex_mask_wf, [('out', 'inputnode.source_files')]),
-                (ds_cortex_mask_wf, merge_cortex_masks, [
-                    ('outputnode.mask_file', f'in{i_hemi + 1}'),
-                ]),
-            ])  # fmt:skip
+        ds_cortex_masks_wf = init_ds_surface_masks_wf(
+            output_dir=output_dir,
+            mask_type='cortex',
+            name='ds_cortex_masks_wf',
+            entities={'extension': '.label.gii'},
+        )
+        workflow.connect([
+            (cortex_masks_wf, ds_cortex_masks_wf, [
+                ('outputnode.cortex_masks', 'inputnode.mask_files'),
+                ('outputnode.source_files', 'inputnode.source_files'),
+            ]),
+            (ds_cortex_masks_wf, outputnode, [('outputnode.mask_files', 'cortex_mask')]),
+        ])  # fmt:skip
     else:
         LOGGER.info('ANAT Stage 11: Found pre-computed cortical surface mask')
         outputnode.inputs.cortex_mask = sorted(precomputed['cortex_mask'])

src/smriprep/workflows/outputs.py

@@ -1231,6 +1231,85 @@ def init_template_iterator_wf(
     return workflow
 
 
+def init_ds_surface_masks_wf(
+    *,
+    output_dir: str,
+    mask_type: ty.Literal['cortex', 'roi', 'ribbon', 'brain'],
+    entities: dict[str, str] | None = None,
+    name='ds_surface_masks_wf',
+) -> Workflow:
+    """Save GIFTI surface masks.
+
+    Parameters
+    ----------
+    output_dir : :class:`str`
+        Directory in which to save derivatives
+    mask_type : :class:`str`
+        Type of mask to save
+    entities : :class:`dict` of :class:`str`
+        Entities to include in outputs
+    name : :class:`str`
+        Workflow name (default: ds_surface_masks_wf)
+
+    Inputs
+    ------
+    source_files : list of lists of str
+        List of lists of source files.
+        Left hemisphere sources first, then right hemisphere sources.
+    mask_files : list of str
+        List of input mask files.
+        Left hemisphere mask first, then right hemisphere mask.
+
+    Outputs
+    -------
+    mask_files : list of str
+        List of output mask files.
+        Left hemisphere mask first, then right hemisphere mask.
+    """
+    workflow = Workflow(name=name)
+
+    if entities is None:
+        entities = {}
+
+    inputnode = pe.Node(
+        niu.IdentityInterface(fields=['mask_files', 'source_files']),
+        name='inputnode',
+    )
+    outputnode = pe.Node(niu.IdentityInterface(fields=['mask_files']), name='outputnode')
+
+    sources = pe.MapNode(
+        niu.Function(function=_bids_relative),
+        name='sources',
+        iterfield='in_files',
+    )
+    sources.inputs.bids_root = output_dir
+
+    ds_mask = pe.MapNode(
+        DerivativesDataSink(
+            base_directory=output_dir,
+            hemi=['L', 'R'],
+            desc=mask_type,
+            **entities,
+        ),
+        iterfield=('in_file', 'hemi', 'source_file'),
+        name='ds_mask',
+        run_without_submitting=True,
+    )
+    if mask_type == 'brain':
+        ds_mask.inputs.Type = 'Brain'
+    else:
+        ds_mask.inputs.Type = 'ROI'
+
+    workflow.connect([
+        (inputnode, ds_mask, [('mask_files', 'in_file')]),
+        (inputnode, sources, [('source_files', 'in_files')]),
+        (sources, ds_mask, [('out', 'source_file')]),
+        (ds_mask, outputnode, [('out_file', 'mask_files')]),
+    ])  # fmt:skip
+
+    return workflow
+
+
 def _bids_relative(in_files, bids_root):
     from pathlib import Path
 

src/smriprep/workflows/surfaces.py

@@ -1156,9 +1156,9 @@ def init_hcp_morphometrics_wf(
     return workflow
 
 
-def init_cortex_mask_wf(
+def init_cortex_masks_wf(
     *,
-    name: str = 'cortex_mask_wf',
+    name: str = 'cortex_masks_wf',
 ):
     """Create a cortical surface mask from a surface file.
 
@@ -1167,55 +1167,100 @@ def init_cortex_mask_wf(
             :graph2use: orig
             :simple_form: yes
 
-            from smriprep.workflows.surfaces import init_cortex_mask_wf
-            wf = init_cortex_mask_wf()
+            from smriprep.workflows.surfaces import init_cortex_masks_wf
+            wf = init_cortex_masks_wf()
 
     Inputs
     ------
-    midthickness : str
-        One hemisphere's FreeSurfer midthickness surface file in GIFTI format
-    thickness : str
-        One hemisphere's FreeSurfer thickness file in GIFTI format
-    hemi : {'L', 'R'}
-        Hemisphere indicator
+    midthickness : list of str
+        Each hemisphere's FreeSurfer midthickness surface file in GIFTI format
+    thickness : list of str
+        Each hemisphere's FreeSurfer thickness file in GIFTI format
 
     Outputs
     -------
-    roi : str
-        Cortical surface mask in GIFTI format
+    cortex_masks : list of str
+        Cortical surface mask in GIFTI format for each hemisphere
     """
     DEFAULT_MEMORY_MIN_GB = 0.01
 
     workflow = Workflow(name=name)
 
     inputnode = pe.Node(
-        niu.IdentityInterface(fields=['midthickness', 'thickness', 'hemi']),
+        niu.IdentityInterface(fields=['midthickness', 'thickness']),
         name='inputnode',
     )
-    outputnode = pe.Node(niu.IdentityInterface(fields=['roi']), name='outputnode')
-
-    # Thickness is presumably already positive, but HCP uses abs(-thickness)
-    abs_thickness = pe.Node(MetricMath(metric='thickness', operation='abs'), name='abs_thickness')
-
-    # Native ROI is thickness > 0, with holes and islands filled
-    initial_roi = pe.Node(MetricMath(metric='roi', operation='bin'), name='initial_roi')
-    fill_holes = pe.Node(MetricFillHoles(), name='fill_holes', mem_gb=DEFAULT_MEMORY_MIN_GB)
-    native_roi = pe.Node(MetricRemoveIslands(), name='native_roi', mem_gb=DEFAULT_MEMORY_MIN_GB)
+    outputnode = pe.Node(niu.IdentityInterface(fields=['cortex_masks']), name='outputnode')
 
+    # Combine the inputs into a list
+    combine_sources = pe.Node(
+        niu.Merge(2, no_flatten=True),
+        name='combine_sources',
+    )
     workflow.connect([
-        (inputnode, abs_thickness, [
-            ('hemi', 'hemisphere'),
-            ('thickness', 'metric_file'),
+        (inputnode, combine_sources, [
+            ('midthickness', 'in1'),
+            ('thickness', 'in2'),
         ]),
-        (inputnode, initial_roi, [('hemi', 'hemisphere')]),
-        (abs_thickness, initial_roi, [('metric_file', 'metric_file')]),
-        (inputnode, fill_holes, [('midthickness', 'surface_file')]),
-        (inputnode, native_roi, [('midthickness', 'surface_file')]),
-        (initial_roi, fill_holes, [('metric_file', 'metric_file')]),
-        (fill_holes, native_roi, [('out_file', 'metric_file')]),
-        (native_roi, outputnode, [('out_file', 'roi')]),
+        (combine_sources, outputnode, [(('out', _transpose_lol), 'source_files')]),
     ])  # fmt:skip
 
+    combine_masks = pe.Node(
+        niu.Merge(2),
+        name='combine_masks',
+    )
+    workflow.connect([(combine_masks, outputnode, [('out', 'cortex_masks')])])
+
+    for i_hemi, hemi in enumerate(['L', 'R']):
+        select_midthickness = pe.Node(
+            niu.Select(index=i_hemi),
+            name=f'select_midthickness_{hemi}',
+        )
+        select_thickness = pe.Node(
+            niu.Select(index=i_hemi),
+            name=f'select_thickness_{hemi}',
+        )
+        workflow.connect([
+            (inputnode, select_midthickness, [('midthickness', 'inlist')]),
+            (inputnode, select_thickness, [('thickness', 'inlist')]),
+        ])  # fmt:skip
+
+         # Thickness is presumably already positive, but HCP uses abs(-thickness)
+        abs_thickness = pe.Node(
+            MetricMath(metric='thickness', operation='abs'),
+            name=f'abs_thickness_{hemi}',
+        )
+
+        # Native ROI is thickness > 0, with holes and islands filled
+        initial_roi = pe.Node(
+            MetricMath(metric='roi', operation='bin'),
+            name=f'initial_roi_{hemi}',
+        )
+        fill_holes = pe.Node(
+            MetricFillHoles(),
+            name=f'fill_holes_{hemi}',
+            mem_gb=DEFAULT_MEMORY_MIN_GB,
+        )
+        native_roi = pe.Node(
+            MetricRemoveIslands(),
+            name=f'native_roi_{hemi}',
+            mem_gb=DEFAULT_MEMORY_MIN_GB,
+        )
+
+        workflow.connect([
+            (inputnode, abs_thickness, [
+                ('hemi', 'hemisphere'),
+                ('thickness', 'metric_file'),
+            ]),
+            (inputnode, initial_roi, [('hemi', 'hemisphere')]),
+            (abs_thickness, initial_roi, [('metric_file', 'metric_file')]),
+            (inputnode, fill_holes, [('midthickness', 'surface_file')]),
+            (inputnode, native_roi, [('midthickness', 'surface_file')]),
+            (initial_roi, fill_holes, [('metric_file', 'metric_file')]),
+            (fill_holes, native_roi, [('out_file', 'metric_file')]),
+            (native_roi, combine_masks, [('out_file', f'in{i_hemi + 1}')]),
+        ])  # fmt:skip
+
     return workflow
 
 
@@ -1758,3 +1803,8 @@ def _select_seg(in_files, segmentation):
 
 def _repeat(seq: list, count: int) -> list:
     return seq * count
+
+
+def _transpose_lol(inlist):
+    """Transpose a list of lists."""
+    return list(map(list, zip(*inlist, strict=False)))