ENH: Separate out inline functions into dedicated interfaces

Author: mgxd

fmriprep/interfaces/bids.py

@@ -5,6 +5,7 @@
 from bids.utils import listify
 from nipype.interfaces.base import (
     DynamicTraitedSpec,
+    File,
     SimpleInterface,
     TraitedSpec,
     isdefined,
@@ -60,3 +61,130 @@ def _run_interface(self, runtime):
         self._results['out'] = out
 
         return runtime
+
+
+class _BIDSSourceFileInputSpec(TraitedSpec):
+    bids_info = traits.Dict(
+        mandatory=True,
+        desc='BIDS information dictionary',
+    )
+    precomputed = traits.Dict({}, usedefault=True, desc='Precomputed BIDS information')
+    sessionwise = traits.Bool(False, usedefault=True, desc='Keep session information')
+    anat_type = traits.Enum('t1w', 't2w', usedefault=True, desc='Anatomical reference type')
+
+
+class _BIDSSourceFileOutputSpec(TraitedSpec):
+    source_file = File(desc='Source file')
+
+
+class BIDSSourceFile(SimpleInterface):
+    input_spec = _BIDSSourceFileInputSpec
+    output_spec = _BIDSSourceFileOutputSpec
+
+    def _run_interface(self, runtime):
+        src = self.inputs.bids_info[self.inputs.anat_type]
+
+        if not src and self.inputs.precomputed.get(f'{self.inputs.anat_type}_preproc'):
+            src = self.inputs.bids_info['bold']
+            self._results['source_file'] = _create_multi_source_file(src)
+            return runtime
+
+        self._results['source_file'] = _create_multi_source_file(
+            src,
+            sessionwise=self.inputs.sessionwise,
+        )
+        return runtime
+
+
+class _CreateFreeSurferIDInputSpec(TraitedSpec):
+    subject_id = traits.Str(mandatory=True, desc='BIDS Subject ID')
+    session_id = traits.Str(desc='BIDS session ID')
+
+
+class _CreateFreeSurferIDOutputSpec(TraitedSpec):
+    subject_id = traits.Str(desc='FreeSurfer subject ID')
+
+
+class CreateFreeSurferID(SimpleInterface):
+    input_spec = _CreateFreeSurferIDInputSpec
+    output_spec = _CreateFreeSurferIDOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results['subject_id'] = _create_fs_id(
+            self.inputs.subject_id,
+            self.inputs.session_id or None,
+        )
+        return runtime
+
+
+def _create_multi_source_file(in_files, sessionwise=False):
+    """
+    Create a generic source name from multiple input files.
+
+    If sessionwise is True, session information from the first file is retained in the name.
+
+    Examples
+    --------
+    >>> _create_multi_source_name([
+    ...     '/path/to/sub-045_ses-test_T1w.nii.gz',
+    ...     '/path/to/sub-045_ses-retest_T1w.nii.gz'])
+    '/path/to/sub-045_T1w.nii.gz'
+    >>> _create_multi_source_name([
+    ...     '/path/to/sub-045_ses-1_run-1_T1w.nii.gz',
+    ...     '/path/to/sub-045_ses-1_run-2_T1w.nii.gz'],
+    ...     sessionwise=True)
+    '/path/to/sub-045_ses-1_T1w.nii.gz'
+    """
+    import re
+    from pathlib import Path
+
+    from nipype.utils.filemanip import filename_to_list
+
+    if not isinstance(in_files, tuple | list):
+        return in_files
+    elif len(in_files) == 1:
+        return in_files[0]
+
+    p = Path(filename_to_list(in_files)[0])
+    try:
+        subj = re.search(r'(?<=^sub-)[a-zA-Z0-9]*', p.name).group()
+        suffix = re.search(r'(?<=_)\w+(?=\.)', p.name).group()
+    except AttributeError as e:
+        raise AttributeError('Could not extract BIDS information') from e
+
+    prefix = f'sub-{subj}'
+
+    if sessionwise:
+        ses = re.search(r'(?<=_ses-)[a-zA-Z0-9]*', p.name)
+        if ses:
+            prefix += f'_ses-{ses.group()}'
+    return str(p.parent / f'{prefix}_{suffix}.nii.gz')
+
+
+def _create_fs_id(subject_id, session_id=None):
+    """
+    Create FreeSurfer subject ID.
+
+    Examples
+    --------
+    >>> _create_fs_id('01')
+    'sub-01'
+    >>> _create_fs_id('sub-01')
+    'sub-01'
+    >>> _create_fs_id('01', 'pre')
+    'sub-01_ses-pre
+    """
+
+    if not subject_id.startswith('sub-'):
+        subject_id = f'sub-{subject_id}'
+
+    if session_id:
+        ses_str = session_id
+        if isinstance(session_id, list):
+            from smriprep.utils.misc import stringify_sessions
+
+            ses_str = stringify_sessions(session_id)
+        if not ses_str.startswith('ses-'):
+            ses_str = f'ses-{ses_str}'
+        subject_id += f'_{ses_str}'
+    return subject_id

fmriprep/workflows/base.py

@@ -176,7 +176,6 @@ def init_single_subject_wf(
     from niworkflows.interfaces.nilearn import NILEARN_VERSION
     from niworkflows.interfaces.utility import KeySelect
     from niworkflows.utils.bids import collect_data
-    from niworkflows.utils.misc import fix_multi_T1w_source_name
     from niworkflows.utils.spaces import Reference
     from smriprep.workflows.anatomical import init_anat_fit_wf
     from smriprep.workflows.outputs import (
@@ -193,6 +192,7 @@ def init_single_subject_wf(
     )
 
     from fmriprep.workflows.bold.base import init_bold_wf
+    from fmriprep.interfaces.bids import BIDSSourceFile, CreateFreeSurferID
 
     if name is None:
         name = f'sub_{subject_id}_wf'
@@ -300,10 +300,20 @@ def init_single_subject_wf(
         name='bidssrc',
     )
 
+    src_file = pe.Node(
+        BIDSSourceFile(
+            precomputed=anatomical_cache,
+            sessionwise=config.workflow.subject_anatomical_reference == 'sessionwise',
+        ),
+        name='source_anatomical',
+    )
+
     bids_info = pe.Node(
         BIDSInfo(bids_dir=config.execution.bids_dir, bids_validate=False), name='bids_info'
     )
 
+    create_fs_id = pe.Node(CreateFreeSurferID(), name='create_fs_id')
+
     summary = pe.Node(
         SubjectSummary(
             std_spaces=spaces.get_spaces(nonstandard=False),
@@ -372,17 +382,15 @@ def init_single_subject_wf(
             'No T1w image found; using precomputed T1w image: %s', anatomical_cache['t1w_preproc']
         )
         workflow.connect([
-            (bidssrc, bids_info, [(('bold', fix_multi_T1w_source_name), 'in_file')]),
             (anat_fit_wf, summary, [('outputnode.t1w_preproc', 't1w')]),
             (anat_fit_wf, ds_report_summary, [('outputnode.t1w_preproc', 'source_file')]),
             (anat_fit_wf, ds_report_about, [('outputnode.t1w_preproc', 'source_file')]),
         ])  # fmt:skip
     else:
         workflow.connect([
-            (bidssrc, bids_info, [(('t1w', fix_multi_T1w_source_name), 'in_file')]),
-            (bidssrc, summary, [('t1w', 't1w')]),
-            (bidssrc, ds_report_summary, [(('t1w', fix_multi_T1w_source_name), 'source_file')]),
-            (bidssrc, ds_report_about, [(('t1w', fix_multi_T1w_source_name), 'source_file')]),
+            (src_file, summary, [('source_file', 't1w')]),
+            (src_file, ds_report_summary, [('source_file', 'source_file')]),
+            (src_file, ds_report_about, [('source_file', 'source_file')]),
         ])  # fmt:skip
 
     workflow.connect([
@@ -393,7 +401,13 @@ def init_single_subject_wf(
             ('roi', 'inputnode.roi'),
             ('flair', 'inputnode.flair'),
         ]),
-        (bids_info, anat_fit_wf, [(('subject', _prefix, 'session'), 'inputnode.subject_id')]),
+        (bidssrc, src_file, [('out_dict', 'bids_info')]),
+        (src_file, bids_info, [('source_file', 'in_file')]),
+        (bids_info, create_fs_id, [
+            ('subject', 'subject_id'),
+            ('session', 'session_id'),
+        ]),
+        (create_fs_id, anat_fit_wf, [('subject_id', 'inputnode.subject_id')]),
         # Reporting connections
         (inputnode, summary, [('subjects_dir', 'subjects_dir')]),
         (bidssrc, summary, [('t2w', 't2w'), ('bold', 'bold')]),
@@ -959,23 +973,6 @@ def map_fieldmap_estimation(
     return fmap_estimators, estimator_map
 
 
-def _prefix(subject_id, session_id=None):
-    """Create FreeSurfer subject ID."""
-    if not subject_id.startswith('sub-'):
-        subject_id = f'sub-{subject_id}'
-
-    if session_id:
-        ses_str = session_id
-        if isinstance(session_id, list):
-            from smriprep.utils.misc import stringify_sessions
-
-            ses_str = stringify_sessions(session_id)
-        if not ses_str.startswith('ses-'):
-            ses_str = f'ses-{ses_str}'
-        subject_id += f'_{ses_str}'
-    return subject_id
-
-
 def clean_datasinks(workflow: pe.Workflow) -> pe.Workflow:
     # Overwrite ``out_path_base`` of smriprep's DataSinks
     for node in workflow.list_node_names():