Merge pull request #134 from nipreps/speedup_hmc

mnoergaard · web-flow · commit 646083a04a9a · 2025-09-16T13:04:40.000+02:00
ENH: Speed up head motion correction workflow
diff --git a/docs/usage.rst b/docs/usage.rst
@@ -192,9 +192,19 @@ registration (mri_robust_register) algorithm utilized settings optimized for PET
 scaling was enabled, automated sensitivity detection was activated, and the Frobenius norm threshold 
 for convergence was set at 0.0001, ensuring precise and consistent alignment across frames.
 
-To edit the motion correction parameters and run the workflow, use: ::
+By default, *PETPrep* evaluates the frames acquired after
+:option:`--hmc-start-time` and initializes motion correction with the
+frame exhibiting the highest tracer uptake. Provide a zero-based index
+with :option:`--hmc-init-frame` to override this choice. Adding
+:option:`--hmc-init-frame-fix` keeps whichever frame is selected (automatic or
+manual) fixed during robust template estimation to improve reproducibility.
+Iterations are automatically disabled to reduce runtime when :option:`--hmc-init-frame-fix` is
+used.
+
+Examples: ::
 
     $ petprep /data/bids_root /out participant --hmc-fwhm 8 --hmc-start-time 60
+    $ petprep /data/bids_root /out participant --hmc-init-frame 10 --hmc-init-frame-fix
 
 Segmentation
 ----------------
diff --git a/docs/workflows.rst b/docs/workflows.rst
@@ -391,8 +391,13 @@ parameters for each time-step are passed on to the
 
 The smoothing kernel width and onset of motion estimation can be
 customized via the :option:`--hmc-fwhm` and :option:`--hmc-start-time`
-command line options.  By default a 10 mm FWHM Gaussian is applied and
-estimation begins at 120 s.
+command line options. By default, PETPrep initializes registration with
+the frame showing the highest tracer uptake after this start time. An
+explicit zero-based frame index can be provided with
+:option:`--hmc-init-frame`. Adding :option:`--hmc-init-frame-fix` keeps the chosen
+frame fixed during robust template estimation and disables iterations to
+reduce runtime. A 10 mm FWHM Gaussian is applied and estimation begins at
+120 s unless otherwise specified.
 
 Pre-processed PET in native space
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/petprep/cli/parser.py b/petprep/cli/parser.py
@@ -92,6 +92,10 @@ def _min_one(value, parser):
             raise parser.error("Argument can't be less than one.")
         return value
 
+    def _int_or_auto(value):
+        """Parse an integer value or the special 'auto' keyword."""
+        return 'auto' if value == 'auto' else int(value)
+
     def _to_gb(value):
         scale = {'G': 1, 'T': 10**3, 'M': 1e-3, 'K': 1e-6, 'B': 1e-9}
         digits = ''.join([c for c in value if c.isdigit()])
@@ -536,6 +540,24 @@ def _bids_filter(value, parser):
         type=float,
         help='Time (in seconds) after which head-motion estimation is performed.',
     )
+    g_hmc.add_argument(
+        '--hmc-init-frame',
+        dest='hmc_init_frame',
+        nargs='?',
+        const='auto',
+        default='auto',
+        type=_int_or_auto,
+        help=(
+            "Initial frame index for head-motion estimation; omit or use 'auto' "
+            'to select the frame with highest uptake.'
+        ),
+    )
+    g_hmc.add_argument(
+        '--hmc-init-frame-fix',
+        dest='hmc_fix_frame',
+        action='store_true',
+        help=('Keep the chosen initial reference frame fixed during head-motion estimation.'),
+    )
 
     g_seg = parser.add_argument_group('Segmentation options')
     g_seg.add_argument(
diff --git a/petprep/cli/tests/test_parser.py b/petprep/cli/tests/test_parser.py
@@ -305,3 +305,23 @@ def test_reference_mask_options(tmp_path, minimal_bids, monkeypatch):
     assert config.workflow.ref_mask_name == 'cerebellum'
     assert config.workflow.ref_mask_index == (3, 4)
     _reset_config()
+
+
+def test_hmc_init_frame_parsing(tmp_path):
+    """Ensure --hmc-init-frame accepts optional integers and defaults to auto."""
+    datapath = tmp_path / 'data'
+    outpath = tmp_path / 'out'
+    datapath.mkdir()
+
+    parser = _build_parser()
+    base_args = [str(datapath), str(outpath), 'participant']
+
+    opts = parser.parse_args(base_args)
+    assert opts.hmc_init_frame == 'auto'
+
+    opts = parser.parse_args(base_args + ['--hmc-init-frame'])
+    assert opts.hmc_init_frame == 'auto'
+
+    opts = parser.parse_args(base_args + ['--hmc-init-frame', '3', '--hmc-init-frame-fix'])
+    assert opts.hmc_init_frame == 3
+    assert opts.hmc_fix_frame is True
diff --git a/petprep/config.py b/petprep/config.py
@@ -596,6 +596,10 @@ class workflow(_Config):
     """FWHM for Gaussian smoothing prior to head-motion estimation."""
     hmc_start_time: float = 120.0
     """Time point (in seconds) at which head-motion estimation starts."""
+    hmc_init_frame: int | str | None = 'auto'
+    """Index of initial frame for head-motion estimation ('auto' selects highest uptake)."""
+    hmc_fix_frame: bool = False
+    """Whether to fix the reference frame during head-motion estimation."""
     seg = 'gtm'
     """Segmentation approach ('gtm', 'brainstem', 'thalamicNuclei',
     'hippocampusAmygdala', 'wm', 'raphe', 'limbic')."""
diff --git a/petprep/data/nipreps.json b/petprep/data/nipreps.json
@@ -181,13 +181,13 @@
     "sub-{subject}[/ses-{session}]/{datatype<perf>|perf}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_rec-{reconstruction}][_dir-{direction}][_run-{run}][_space-{space}][_atlas-{atlas}][_cohort-{cohort}][_desc-{desc}]_{suffix<timeseries>}{extension<.json|.tsv>|.tsv}",
     "sub-{subject}[/ses-{session}]/{datatype<perf>|perf}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_rec-{reconstruction}][_dir-{direction}][_run-{run}][_space-{space}][_atlas-{atlas}][_cohort-{cohort}][_desc-{desc}]_{suffix<asl|aslref|att|cbf|coverage|mask>}{extension<.nii|.nii.gz|.json|.tsv>|.tsv}",
     "sub-{subject}[/ses-{session}]/{datatype<fmap>|fmap}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_dir-{direction}][_run-{run}][_part-{part}][_space-{space}][_cohort-{cohort}][_res-{resolution}][_fmapid-{fmapid}][_desc-{desc}]_{suffix<fieldmap>}{extension<.nii|.nii.gz|.json>|.nii.gz}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}][_desc-{desc}]_{suffix<pet|petref>}{extension<.nii|.nii.gz|.json>|.nii.gz}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_hemi-{hemi<L|R>}]_from-{from}_to-{to}_mode-{mode<image|points>|image}_{suffix<xfm>|xfm}{extension<.txt|.h5>}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}]_hemi-{hemi<L|R>}[_space-{space}][_cohort-{cohort}][_den-{density}][_desc-{desc}]_{suffix<white|smoothwm|pial|midthickness|inflated|vinflated|sphere|flat|sulc|curv|thickness>}{extension<.surf.gii|.shape.gii>}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_den-{density}][_desc-{desc}]_{suffix<sulc|curv|thickness>}{extension<.dscalar.nii|.json>}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}]_desc-{desc}_{suffix<mask>|mask}{extension<.nii|.nii.gz|.json>|.nii.gz}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}]_label-{label}[_desc-{desc}]_{suffix<probseg>|probseg}{extension<.nii|.nii.gz|.json>|.nii.gz}",
-    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_dir-{direction}][_run-{run}][_part-{part}][_space-{space}][_atlas-{atlas}][_cohort-{cohort}][_desc-{desc}]_{suffix<timeseries|regressors>|timeseries}{extension<.json|.tsv>|.tsv}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}][_desc-{desc}]_{suffix<pet|petref>}{extension<.nii|.nii.gz|.json>|.nii.gz}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_hemi-{hemi<L|R>}]_from-{from}_to-{to}_mode-{mode<image|points>|image}_{suffix<xfm>|xfm}{extension<.txt|.h5>}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}]_hemi-{hemi<L|R>}[_space-{space}][_cohort-{cohort}][_den-{density}][_desc-{desc}]_{suffix<white|smoothwm|pial|midthickness|inflated|vinflated|sphere|flat|sulc|curv|thickness>}{extension<.surf.gii|.shape.gii>}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_den-{density}][_desc-{desc}]_{suffix<sulc|curv|thickness>}{extension<.dscalar.nii|.json>}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}]_desc-{desc}_{suffix<mask>|mask}{extension<.nii|.nii.gz|.json>|.nii.gz}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_res-{resolution}]_label-{label}[_desc-{desc}]_{suffix<probseg>|probseg}{extension<.nii|.nii.gz|.json>|.nii.gz}",
+    "sub-{subject}[/ses-{session}]/{datatype<pet>|pet}/sub-{subject}[_ses-{session}][_task-{task}][_acq-{acquisition}][_ce-{ceagent}][_trc-{tracer}][_rec-{reconstruction}][_dir-{direction}][_run-{run}][_part-{part}][_space-{space}][_atlas-{atlas}][_cohort-{cohort}][_desc-{desc}]_{suffix<timeseries|regressors>|timeseries}{extension<.json|.tsv>|.tsv}",
     "sub-{subject}/{datatype<figures>}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_desc-{desc}]_{suffix<T1w|T2w|T1rho|T1map|T2map|T2star|FLAIR|FLASH|PDmap|PD|PDT2|inplaneT[12]|angio|dseg|mask|T2starw|MTw|TSE|pet>}{extension<.html|.svg>|.svg}",
     "sub-{subject}/{datatype<figures>}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_rec-{reconstruction}][_run-{run}][_space-{space}][_cohort-{cohort}][_fmapid-{fmapid}][_desc-{desc}]_{suffix<fieldmap>}{extension<.html|.svg>|.svg}",
     "sub-{subject}/{datatype<figures>}/sub-{subject}[_ses-{session}][_acq-{acquisition}][_ce-{ceagent}][_rec-{reconstruction}][_dir-{direction}][_run-{run}][_echo-{echo}][_part-{part}][_space-{space}][_cohort-{cohort}][_desc-{desc}]_{suffix<dwi|epi|epiref|pet>}{extension<.html|.svg>|.svg}",
diff --git a/petprep/workflows/pet/fit.py b/petprep/workflows/pet/fit.py
@@ -281,6 +281,8 @@ def init_pet_fit_wf(
             start_time=config.workflow.hmc_start_time,
             frame_durations=frame_durations,
             frame_start_times=frame_start_times,
+            initial_frame=config.workflow.hmc_init_frame,
+            fixed_frame=config.workflow.hmc_fix_frame,
         )
 
         ds_hmc_wf = init_ds_hmc_wf(
diff --git a/petprep/workflows/pet/hmc.py b/petprep/workflows/pet/hmc.py
@@ -100,6 +100,14 @@ def lta_list(in_file):
     return lta_list
 
 
+def _find_highest_uptake_frame(in_files: list[str]) -> int:
+    import nibabel as nb
+    import numpy as np
+
+    uptake = [np.sum(nb.load(f).get_fdata(dtype=np.float32)) for f in in_files]
+    return int(np.argmax(uptake)) + 1
+
+
 class _LTAList2ITKInputSpec(BaseInterfaceInputSpec):
     in_xforms = InputMultiObject(File(exists=True), mandatory=True)
     in_reference = File(exists=True, mandatory=True)
@@ -139,6 +147,8 @@ def init_pet_hmc_wf(
     start_time: float = 120.0,
     frame_durations: Sequence[float] | None = None,
     frame_start_times: Sequence[float] | None = None,
+    initial_frame: int | str | None = 'auto',
+    fixed_frame: bool = False,
     name: str = 'pet_hmc_wf',
 ):
     r"""
@@ -177,6 +187,15 @@ def init_pet_hmc_wf(
     frame_start_times : :class:`~typing.Sequence`\[:obj:`float`] or ``None``
         Optional list of frame onset times used together with
         ``frame_durations`` to locate the start frame.
+    initial_frame : :obj:`int`, ``'auto'`` or ``None``
+        0-based index of the frame used to initialize motion correction. If ``'auto'`` or
+        ``None`` (default), the frame with the highest uptake is selected automatically.
+        FreeSurfer's ``initial_timepoint`` is 1-based; this workflow applies the
+        required offset internally.
+    fixed_frame : :obj:`bool`
+        Whether to keep the initial time point fixed during robust template
+        estimation (``fs.RobustTemplate``'s ``fixtp`` parameter). If ``True``,
+        iterations are skipped to reduce runtime.
     name : :obj:`str`
         Name of workflow (default: ``pet_hmc_wf``)
 
@@ -274,6 +293,17 @@ def num_files(filelist):
         name='create_lta_list',
     )
 
+    auto_init_frame = initial_frame in (None, 'auto')
+    if auto_init_frame:
+        find_highest_uptake_frame = pe.Node(
+            niu.Function(
+                input_names=['in_files'],
+                output_names=['index'],
+                function=_find_highest_uptake_frame,
+            ),
+            name='find_highest_uptake_frame',
+        )
+
     # Motion estimation
     robust_template = pe.Node(
         fs.RobustTemplate(
@@ -282,9 +312,13 @@ def num_files(filelist):
             average_metric='mean',
             args='--cras',
             num_threads=omp_nthreads,
+            fixed_timepoint=fixed_frame,
+            no_iteration=fixed_frame,
         ),
         name='est_robust_hmc',
     )
+    if not auto_init_frame:
+        robust_template.inputs.initial_timepoint = int(initial_frame) + 1
     upd_xfm = pe.Node(
         niu.Function(
             input_names=['xforms', 'idx'],
@@ -324,4 +358,12 @@ def num_files(filelist):
         (ref_to_nii, outputnode, [('out_file', 'petref')]),
     ])  # fmt:skip
 
+    if auto_init_frame:
+        workflow.connect(
+            [
+                (thresh, find_highest_uptake_frame, [('out_file', 'in_files')]),
+                (find_highest_uptake_frame, robust_template, [('index', 'initial_timepoint')]),
+            ]
+        )
+
     return workflow
diff --git a/petprep/workflows/pet/tests/test_hmc.py b/petprep/workflows/pet/tests/test_hmc.py
@@ -1,4 +1,13 @@
-from ..hmc import get_start_frame, init_pet_hmc_wf, update_list_transforms
+import nibabel as nb
+import numpy as np
+import pytest
+
+from ..hmc import (
+    _find_highest_uptake_frame,
+    get_start_frame,
+    init_pet_hmc_wf,
+    update_list_transforms,
+)
 
 
 def test_get_start_frame_basic():
@@ -26,11 +35,8 @@ def test_update_list_transforms_padding():
     assert update_list_transforms(xforms, 0) == xforms
 
 
-import pytest
-
-
 def test_update_list_transforms_empty():
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match='cannot be empty'):
         update_list_transforms([], 1)
 
 
@@ -40,3 +46,34 @@ def test_init_pet_hmc_wf_nodes():
     assert 'split_frames' in names
     assert 'est_robust_hmc' in names
     assert 'convert_ref' in names
+
+
+def test_init_pet_hmc_wf_auto_inittp():
+    wf = init_pet_hmc_wf(mem_gb=1, omp_nthreads=1, initial_frame='auto')
+    names = wf.list_node_names()
+    assert 'find_highest_uptake_frame' in names
+
+
+def test_init_pet_hmc_wf_specific_inittp():
+    wf = init_pet_hmc_wf(mem_gb=1, omp_nthreads=1, initial_frame=2, fixed_frame=True)
+    names = wf.list_node_names()
+    assert 'find_highest_uptake_frame' not in names
+    node = wf.get_node('est_robust_hmc')
+    initial_frame = 2
+    assert node.inputs.initial_timepoint == initial_frame + 1
+    assert node.inputs.fixed_timepoint is True
+    assert node.inputs.no_iteration is True
+
+
+def test_find_highest_uptake_frame(tmp_path):
+    data = [np.ones((2, 2, 2)) * i for i in (1, 2, 3)]
+    files = []
+    for idx, arr in enumerate(data):
+        img = nb.Nifti1Image(arr, np.eye(4))
+        fname = tmp_path / f'frame{idx}.nii.gz'
+        img.to_filename(fname)
+        files.append(str(fname))
+
+    expected = np.argmax([arr.sum() for arr in data]) + 1
+    result = _find_highest_uptake_frame(files)
+    assert result == expected
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,7 +27,7 @@ dependencies = [
     "nireports >= 24.1.0",
     "nitime >= 0.9",
     "nitransforms >= 24.1.1",
-    "niworkflows >= 1.14.0",
+    "niworkflows @ git+https://github.com/nipreps/niworkflows.git@master",
     "numpy >= 1.24",
     "packaging >= 24",
     "pandas >= 1.2",

Original file line number	Diff line number	Diff line change
`@@ -281,6 +281,8 @@ def init_pet_fit_wf(`
`281`	`281`	`start_time=config.workflow.hmc_start_time,`
`282`	`282`	`frame_durations=frame_durations,`
`283`	`283`	`frame_start_times=frame_start_times,`
	`284`	`+ initial_frame=config.workflow.hmc_init_frame,`
	`285`	`+ fixed_frame=config.workflow.hmc_fix_frame,`
`284`	`286`	`)`
`285`	`287`
`286`	`288`	`ds_hmc_wf = init_ds_hmc_wf(`