Merge pull request #5 from mnoergaard/main

refactor fMRIPrep to PETPrep

Author: mnoergaard

fmriprep/workflows/base.py

@@ -146,7 +146,7 @@ def init_single_subject_wf(subject_id: str):
 
     """
     from niworkflows.engine.workflows import LiterateWorkflow as Workflow
-    from niworkflows.interfaces.bids import BIDSDataGrabber, BIDSInfo
+    from niworkflows.interfaces.bids import BIDSInfo, BIDSDataGrabber
     from niworkflows.interfaces.nilearn import NILEARN_VERSION
     from niworkflows.interfaces.utility import KeySelect
     from niworkflows.utils.bids import collect_data
@@ -209,7 +209,7 @@ def init_single_subject_wf(subject_id: str):
         config.execution.bids_dir,
         subject_id,
         bids_filters=config.execution.bids_filters,
-        queries=queries
+        queries=queries,
     )[0]
 
 
@@ -455,12 +455,12 @@ def init_single_subject_wf(subject_id: str):
             )
             ds_grayord_metrics_wf = init_ds_grayord_metrics_wf(
                 bids_root=bids_root,
-                output_dir=fmriprep_dir,
+                output_dir=petprep_dir,
                 metrics=['curv', 'thickness', 'sulc'],
                 cifti_output=config.workflow.cifti_output,
             )
             ds_fsLR_surfaces_wf = init_ds_surfaces_wf(
-                output_dir=fmriprep_dir,
+                output_dir=petprep_dir,
                 surfaces=['white', 'pial', 'midthickness'],
                 entities={
                     'space': 'fsLR',

fmriprep/workflows/pet/confounds.py

@@ -36,7 +36,6 @@
 from ...config import DEFAULT_MEMORY_MIN_GB
 from ...interfaces import DerivativesDataSink
 from ...interfaces.confounds import (
-    FilterDropped,
     PETSummary,
     FramewiseDisplacement,
     FSLMotionParams,
@@ -148,7 +147,7 @@ def init_pet_confs_wf(
     from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
     from niworkflows.interfaces.images import SignalExtraction
     from niworkflows.interfaces.morphology import BinaryDilation, BinarySubtraction
-    from niworkflows.interfaces.nibabel import ApplyMask, Binarize
+    from niworkflows.interfaces.nibabel import Binarize
     from niworkflows.interfaces.utility import AddTSVHeader, DictMerge
 
     from ...interfaces.confounds import aCompCorMasks
@@ -262,6 +261,18 @@ def init_pet_confs_wf(
         'white_matter',
         'csf_wm',
     ]
+    get_pet_zooms = pe.Node(niu.Function(function=_get_zooms), name='get_pet_zooms')
+    acompcor_masks = pe.Node(aCompCorMasks(), name='acompcor_masks')
+    acompcor_tfm = pe.MapNode(
+        ApplyTransforms(interpolation='MultiLabel', invert_transform_flags=[True]),
+        name='acompcor_tfm',
+        iterfield=['input_image'],
+    )
+    acompcor_bin = pe.MapNode(
+        Binarize(thresh_low=0.99),
+        name='acompcor_bin',
+        iterfield=['in_file'],
+    )
     merge_rois = pe.Node(
         niu.Merge(3, ravel_inputs=True), name='merge_rois', run_without_submitting=True
     )
@@ -372,7 +383,20 @@ def _select_cols(table):
         (subtract_mask, outputnode, [('out_mask', 'crown_mask')]),
         # Global signals extraction (constrained by anatomy)
         (inputnode, signals, [('pet', 'in_file')]),
-        (inputnode, merge_rois, [('pet_mask', 'in1')]),
+        (inputnode, get_pet_zooms, [('pet', 'in_file')]),
+        (inputnode, acompcor_masks, [('t1w_tpms', 'in_vfs')]),
+        (get_pet_zooms, acompcor_masks, [('out', 'pet_zooms')]),
+        (acompcor_masks, acompcor_tfm, [('out_masks', 'input_image')]),
+        (inputnode, acompcor_tfm, [
+            ('pet_mask', 'reference_image'),
+            ('petref2anat_xfm', 'transforms'),
+        ]),
+        (acompcor_tfm, acompcor_bin, [('output_image', 'in_file')]),
+        (acompcor_bin, merge_rois, [
+            (('out_mask', _last), 'in3'),
+            (('out_mask', lambda masks: masks[0]), 'in1'),
+            (('out_mask', lambda masks: masks[1]), 'in2'),
+        ]),
         (merge_rois, signals, [('out', 'label_files')]),
 
         # Collate computed confounds together

fmriprep/workflows/pet/fit.py

@@ -37,6 +37,7 @@
 from .outputs import (
     init_ds_hmc_wf,
     init_ds_petref_wf,
+    init_ds_petmask_wf,
     init_ds_registration_wf,
     init_func_fit_reports_wf,
     prepare_timing_parameters,
@@ -176,7 +177,7 @@ def init_pet_fit_wf(
     workflow.add_nodes([inputnode])
 
     petref_buffer = pe.Node(
-        niu.IdentityInterface(fields=['petref', 'pet_file', 'pet_mask']),
+        niu.IdentityInterface(fields=['petref', 'pet_file']),
         name='petref_buffer',
     )
     hmc_buffer = pe.Node(niu.IdentityInterface(fields=['hmc_xforms']), name='hmc_buffer')
@@ -197,11 +198,7 @@ def init_pet_fit_wf(
 
     summary = pe.Node(
         FunctionalSummary(
-            registration=(
-                'Precomputed'
-                if petref2anat_xform
-                else 'mri_coreg'
-            ),
+            registration=('Precomputed' if petref2anat_xform else 'mri_coreg'),
             registration_dof=config.workflow.pet2anat_dof,
             orientation=orientation,
         ),
@@ -218,7 +215,6 @@ def init_pet_fit_wf(
     workflow.connect([
         (petref_buffer, outputnode, [
             ('petref', 'petref'),
-            ('pet_mask', 'pet_mask'),
         ]),
         (hmc_buffer, outputnode, [
             ('hmc_xforms', 'motion_xfm'),
@@ -252,7 +248,6 @@ def init_pet_fit_wf(
             pet_file=pet_file,
             reference_frame=config.workflow.reference_frame,
         )
-        petref_wf.inputs.inputnode.dummy_scans = config.workflow.dummy_scans
 
         ds_petref_wf = init_ds_petref_wf(
             bids_root=layout.root,
@@ -262,6 +257,21 @@ def init_pet_fit_wf(
         )
         ds_petref_wf.inputs.inputnode.source_files = [pet_file]
 
+        # Ensure all stage-1 workflows were created successfully before
+        # attempting to connect them. Nipype's ``connect`` call will fail
+        # with a ``NoneType`` error if any node is undefined.
+        stage1_nodes = [
+            petref_wf,
+            petref_buffer,
+            ds_petref_wf,
+            func_fit_reports_wf,
+            petref_source_buffer,
+        ]
+        if any(node is None for node in stage1_nodes):
+            raise RuntimeError(
+                'PET reference stage could not be built - check inputs and configuration.'
+            )
+
         workflow.connect([
             (petref_wf, petref_buffer, [
                 ('outputnode.pet_file', 'pet_file'),
@@ -344,6 +354,45 @@ def init_pet_fit_wf(
     else:
         outputnode.inputs.petref2anat_xfm = petref2anat_xform
 
+    # Stage 4: Estimate PET brain mask
+    from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
+    from niworkflows.interfaces.nibabel import Binarize
+
+    from .confounds import _binary_union
+
+    t1w_mask_tfm = pe.Node(
+        ApplyTransforms(interpolation='MultiLabel', invert_transform_flags=[True]),
+        name='t1w_mask_tfm',
+    )
+    petref_mask = pe.Node(Binarize(thresh_low=0.2), name='petref_mask')
+    merge_mask = pe.Node(niu.Function(function=_binary_union), name='merge_mask')
+
+    if not petref2anat_xform:
+        workflow.connect(
+            [(pet_reg_wf, t1w_mask_tfm, [('outputnode.itk_pet_to_t1', 'transforms')])]
+        )
+    else:
+        t1w_mask_tfm.inputs.transforms = petref2anat_xform
+
+    workflow.connect(
+        [
+            (inputnode, t1w_mask_tfm, [('t1w_mask', 'input_image')]),
+            (petref_buffer, t1w_mask_tfm, [('petref', 'reference_image')]),
+            (petref_buffer, petref_mask, [('petref', 'in_file')]),
+            (petref_mask, merge_mask, [('out_mask', 'mask1')]),
+            (t1w_mask_tfm, merge_mask, [('output_image', 'mask2')]),
+            (merge_mask, outputnode, [('out', 'pet_mask')]),
+        ]
+    )
+
+    ds_petmask_wf = init_ds_petmask_wf(
+        output_dir=config.execution.petprep_dir,
+        desc='brain',
+        name='ds_petmask_wf',
+    )
+    ds_petmask_wf.inputs.inputnode.source_files = [pet_file]
+    workflow.connect([(merge_mask, ds_petmask_wf, [('out', 'inputnode.petmask')])])
+
     return workflow
 
 
@@ -436,18 +485,13 @@ def init_pet_native_wf(
     )
     outputnode.inputs.metadata = metadata
 
-    petbuffer = pe.Node(
-        niu.IdentityInterface(fields=['pet_file']), name='petbuffer'
-    )
+    petbuffer = pe.Node(niu.IdentityInterface(fields=['pet_file']), name='petbuffer')
 
     # PET source: track original PET file(s)
     # The Select interface requires an index to choose from ``inlist``. Since
     # ``pet_file`` is a single path, explicitly set the index to ``0`` to avoid
     # missing mandatory input errors when the node runs.
-    pet_source = pe.Node(
-        niu.Select(inlist=[pet_file], index=0),
-        name='pet_source'
-    )
+    pet_source = pe.Node(niu.Select(inlist=[pet_file], index=0), name='pet_source')
     validate_pet = pe.Node(ValidateImage(), name='validate_pet')
     workflow.connect([
         (pet_source, validate_pet, [('out', 'in_file')]),

fmriprep/workflows/pet/reference.py

@@ -137,6 +137,8 @@ def init_raw_petref_wf(
             ]
         )  # fmt:skip
 
+    return workflow
+
 
 def init_validation_and_dummies_wf(
     pet_file=None,

fmriprep/workflows/pet/tests/test_confounds.py

@@ -0,0 +1,34 @@
+import nibabel as nb
+import numpy as np
+
+from ..confounds import init_pet_confs_wf
+
+
+def test_dvars_connects_pet_mask(tmp_path):
+    """Check dvars node connection and execution."""
+    wf = init_pet_confs_wf(
+        mem_gb=0.01,
+        metadata={},
+        regressors_all_comps=False,
+        regressors_dvars_th=1.5,
+        regressors_fd_th=0.5,
+    )
+
+    edge = wf._graph.get_edge_data(wf.get_node('inputnode'), wf.get_node('dvars'))
+    assert ('pet_mask', 'in_mask') in edge['connect']
+
+    img = nb.Nifti1Image(np.random.rand(2, 2, 2, 5), np.eye(4))
+    mask = nb.Nifti1Image(np.ones((2, 2, 2), dtype=np.uint8), np.eye(4))
+    pet_file = tmp_path / 'pet.nii.gz'
+    mask_file = tmp_path / 'mask.nii.gz'
+    img.to_filename(pet_file)
+    mask.to_filename(mask_file)
+
+    node = wf.get_node('dvars')
+    node.base_dir = tmp_path
+    node.inputs.in_file = str(pet_file)
+    node.inputs.in_mask = str(mask_file)
+    result = node.run()
+
+    assert result.outputs.out_nstd
+    assert result.outputs.out_std

fmriprep/workflows/pet/tests/test_fit.py

@@ -134,3 +134,23 @@ def test_pet_native_precomputes(
 
     flatgraph = wf._create_flat_graph()
     generate_expanded_graph(flatgraph)
+
+
+def test_pet_fit_mask_connections(bids_root: Path, tmp_path: Path):
+    """Ensure the PET mask is generated and connected correctly."""
+    pet_file = str(bids_root / 'sub-01' / 'pet' / 'sub-01_task-rest_run-1_pet.nii.gz')
+    img = nb.Nifti1Image(np.zeros((2, 2, 2, 1)), np.eye(4))
+    img.to_filename(pet_file)
+
+    with mock_config(bids_dir=bids_root):
+        wf = init_pet_fit_wf(pet_file=pet_file, precomputed={}, omp_nthreads=1)
+
+    assert 'merge_mask' in wf.list_node_names()
+    assert 'ds_petmask_wf.ds_petmask' in wf.list_node_names()
+
+    merge_mask = wf.get_node('merge_mask')
+    edge = wf._graph.get_edge_data(merge_mask, wf.get_node('outputnode'))
+    assert ('out', 'pet_mask') in edge['connect']
+
+    ds_edge = wf._graph.get_edge_data(merge_mask, wf.get_node('ds_petmask_wf'))
+    assert ('out', 'inputnode.petmask') in ds_edge['connect']

fmriprep/workflows/pet/tests/test_pet_mask.py

@@ -0,0 +1,29 @@
+import nibabel as nb
+import numpy as np
+from niworkflows.utils.testing import generate_bids_skeleton
+
+from ...tests import mock_config
+from ...tests.test_base import BASE_LAYOUT
+from ..base import init_pet_wf
+
+
+def test_pet_mask_flow(tmp_path):
+    bids_dir = tmp_path / 'bids'
+    generate_bids_skeleton(bids_dir, BASE_LAYOUT)
+    img = nb.Nifti1Image(np.zeros((2, 2, 2, 10)), np.eye(4))
+    pet_file = bids_dir / 'sub-01' / 'pet' / 'sub-01_task-rest_run-1_pet.nii.gz'
+    img.to_filename(pet_file)
+
+    with mock_config(bids_dir=bids_dir):
+        wf = init_pet_wf(pet_series=str(pet_file), precomputed={})
+
+    edge = wf._graph.get_edge_data(
+        wf.get_node('pet_fit_wf'), wf.get_node('pet_confounds_wf')
+    )
+    assert ('pet_mask', 'inputnode.pet_mask') in edge['connect']
+
+    conf_wf = wf.get_node('pet_confounds_wf')
+    conf_edge = conf_wf._graph.get_edge_data(
+        conf_wf.get_node('inputnode'), conf_wf.get_node('dvars')
+    )
+    assert ('pet_mask', 'in_mask') in conf_edge['connect']

fmriprep/workflows/pet/tests/test_reference.py

@@ -6,22 +6,22 @@
 
 def test_reference_frame_select(tmp_path):
     img = nb.Nifti1Image(np.zeros((5, 5, 5, 4)), np.eye(4))
-    pet_file = tmp_path / "pet.nii.gz"
+    pet_file = tmp_path / 'pet.nii.gz'
     img.to_filename(pet_file)
 
     wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame=2)
     node_names = [n.name for n in wf._get_all_nodes()]
-    assert "extract_frame" in node_names
-    assert "gen_avg" not in node_names
-    node = wf.get_node("extract_frame")
+    assert 'extract_frame' in node_names
+    assert 'gen_avg' not in node_names
+    node = wf.get_node('extract_frame')
     assert node.interface.inputs.t_min == 2
 
 
 def test_reference_frame_average(tmp_path):
     img = nb.Nifti1Image(np.zeros((5, 5, 5, 4)), np.eye(4))
-    pet_file = tmp_path / "pet.nii.gz"
+    pet_file = tmp_path / 'pet.nii.gz'
     img.to_filename(pet_file)
 
-    wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame="average")
+    wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame='average')
     node_names = [n.name for n in wf._get_all_nodes()]
-    assert "gen_avg" in node_names
+    assert 'gen_avg' in node_names