ENH: Optimize co-registration of T1w to petref, and provide more options for petref generation

docs/usage.rst

@@ -211,16 +211,37 @@ Examples: ::
     $ petprep /data/bids_root /out participant --hmc-init-frame 10 --hmc-init-frame-fix
     $ petprep /data/bids_root /out participant --hmc-off
 
+
+Reference image selection
+-------------------------
+Use :option:`--petref` to control how the reference volume is built from the
+dynamic PET series. Each strategy uses the frame timing metadata from
+``FrameTimesStart`` and ``FrameDuration`` to weight volumes; missing metadata
+will raise an error before preprocessing starts.
+
+* ``template`` (default) reuses the motion-correction template, providing a
+  consistent target for downstream registration. When :option:`--hmc-off`
+  disables motion correction, requesting ``template`` automatically falls back
+  to ``twa`` with a warning.
+* ``twa`` computes a time-weighted average, which often emphasizes later frames with
+  higher counts and longer durations.
+* ``sum`` produces a straightforward summed image.
+* ``first5min`` averages only the first 5 minutes of PET data to capture perfusion-like
+  uptake.
+
 Anatomical co-registration
 --------------------------
 *PETPrep* aligns the PET reference volume to the T1-weighted anatomy before
-deriving downstream outputs. By default, FreeSurfer's ``mri_coreg`` performs
-the alignment, with the :option:`--pet2anat-dof` flag controlling the degrees
-of freedom (rigid-body, 6 dof, is the default). When working with low
-signal-to-noise references or challenging anatomy, the
-:option:`--pet2anat-robust` flag enables ``mri_robust_register`` with an NMI
-cost function to improve robustness. This mode is restricted to rigid-body
-alignment and therefore requires ``--pet2anat-dof 6``.
+deriving downstream outputs. The anatomical image is first trimmed with
+FSL's ``robustfov`` to remove the shoulder/neck and masked to limit registration to brain voxels. Choose
+the registration backend with :option:`--pet2anat-method`: ``mri_coreg``
+(default FreeSurfer co-registration), ``robust`` (FreeSurfer
+``mri_robust_register`` with an NMI cost function), or ``ants`` (ANTs rigid
+registration that consumes the unmasked T1w and a separate mask). The
+:option:`--pet2anat-dof` flag controls the degrees of freedom; ``robust`` and
+``ants`` are limited to rigid-body alignment and therefore require
+``--pet2anat-dof 6``. All modes emit paired ITK transforms for reuse in later
+resampling steps.
 
 Segmentation
 ----------------

docs/workflows.rst

@@ -342,29 +342,17 @@ PET reference image estimation
     wf = init_raw_petref_wf()
 
 This workflow estimates a reference image for a
-:abbr:`PET (positron emission tomography)` series as follows:
-When T1-saturation effects ("dummy scans" or non-steady state volumes) are
-detected, they are averaged and used as reference due to their
-superior tissue contrast.
-Otherwise, a median of motion corrected subset of volumes is used.
-
-This reference is used for :ref:`head-motion estimation <pet_hmc>`.
-
-For the :ref:`registration workflow <pet_reg>`, the reference image is
-either the above described reference image or a single-band reference,
-if one is found in the input dataset.
-In either case, this image is contrast-enhanced and skull-stripped
-(see :py:func:`~niworkflows.func.util.init_enhance_and_skullstrip_bold_wf`).
-If fieldmaps are present, the skull-stripped reference is corrected
-prior to registration.
-
-.. figure:: _static/sub-01_task-balloonanalogrisktask_run-1_desc-rois_bold.svg
-
-    The red contour shows the brain mask estimated for a PET reference volume.
-    The blue and magenta contours show the tCompCor and aCompCor masks,
-    respectively. (See :ref:`pet_confounds`, below.)
-
-.. _pet_hmc:
+:abbr:`PET (positron emission tomography)` series according to the strategy
+requested with :option:`--petref`:
+
+* ``template`` (default) uses the motion-correction template built during
+  head-motion estimation. When head-motion correction is disabled, ``template``
+  automatically falls back to ``twa`` while emitting a warning.
+* ``twa`` computes a time-weighted average of the motion-corrected series,
+  preserving dynamic information when later frames carry more counts.
+* ``sum`` produces a summed image of the motion-corrected series.
+* ``first5min`` weights only the first five minutes of the acquisition, which
+  can be helpful for tracers whose early dynamics resemble perfusion.
 
 Head-motion estimation
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -438,17 +426,19 @@ PET to T1w registration
         omp_nthreads=1,
         pet2anat_dof=6,
         mem_gb=3,
-        use_robust_register=False,
+        pet2anat_method='mri_coreg',
     )
 
 The PET reference volume is aligned to the skull-stripped anatomical image
-using FreeSurfer's ``mri_coreg`` with the number of degrees of freedom set via
-the :option:`--pet2anat-dof` flag. The resulting affine is converted to ITK
-format for downstream application, along with its inverse.
-
-If co-registration proves challenging, the :option:`--pet2anat-robust` flag
-switches the workflow to FreeSurfer's ``mri_robust_register`` with an NMI cost function and restricted
-to rigid-body (6 dof) transforms. This method is more robust to large initial misalignments.
+using the method selected via :option:`--pet2anat-method`. The anatomical image
+is first cropped with FSL's ``robustfov`` and masked to focus the alignment on
+brain tissue (ANTs receives the unmasked cropped image together with its mask).
+By default, the workflow runs FreeSurfer's ``mri_coreg`` with the number of
+degrees of freedom set via the :option:`--pet2anat-dof` flag. Alternative modes
+include FreeSurfer's ``mri_robust_register`` (``--pet2anat-method robust``) and
+ANTs rigid registration (``--pet2anat-method ants``). Both alternatives are
+limited to rigid-body alignment (6 DoF). The resulting affine is converted to
+ITK format for downstream application, along with its inverse.
 
 Resampling PET runs onto standard spaces
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -464,11 +454,10 @@ Resampling PET runs onto standard spaces
     )
 
 This sub-workflow concatenates the transforms calculated upstream (see
-`Head-motion estimation`_, `Susceptibility Distortion Correction (SDC)`_ --if
-fieldmaps are available--, `EPI to T1w registration`_, and an anatomical-to-standard
+`Head-motion estimation`_, `PET to T1w registration`_, and an anatomical-to-standard
 transform from `Preprocessing of structural MRI`_) to map the
-:abbr:`EPI (echo-planar imaging)`
-image to the standard spaces given by the ``--output-spaces`` argument
+:abbr:`PET (positron emission tomography)` series from native space to
+the standard spaces given by the ``--output-spaces`` argument
 (see :ref:`output-spaces`).
 It also maps the T1w-based mask to each of those standard spaces.
 
@@ -478,7 +467,7 @@ step, so as little information is lost as possible.
 The output space grid can be specified using modifiers to the ``--output-spaces``
 argument.
 
-EPI sampled to FreeSurfer surfaces
+PET sampled to FreeSurfer surfaces
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 :py:func:`~petprep.workflows.pet.resampling.init_pet_surf_wf`
 
@@ -495,7 +484,7 @@ EPI sampled to FreeSurfer surfaces
         output_dir='.',
     )
 
-If FreeSurfer processing is enabled, the motion-corrected functional series
+If FreeSurfer processing is enabled, the motion-corrected PET series
 (after single shot resampling to T1w space) is sampled to the
 surface by averaging across the cortical ribbon.
 Specifically, at each vertex, the segment normal to the white-matter surface, extending to the pial

petprep/cli/parser.py

@@ -358,10 +358,14 @@ def _bids_filter(value, parser):
         '6 degrees (rotation and translation) are used by default.',
     )
     g_conf.add_argument(
-        '--pet2anat-robust',
-        action='store_true',
-        help='Use FreeSurfer mri_robust_register with an NMI cost function for'
-        'PET-to-T1w co-registration. This option is limited to 6 dof.',
+        '--pet2anat-method',
+        action='store',
+        default='mri_coreg',
+        choices=['mri_coreg', 'robust', 'ants'],
+        help='Method for PET-to-anatomical registration. '
+        '"mri_coreg" (default) uses FreeSurfer mri_coreg. '
+        '"robust" uses FreeSurfer mri_robust_register (6 DoF only). '
+        '"ants" uses ANTs rigid registration (6 DoF only).',
     )
     g_conf.add_argument(
         '--force-bbr',
@@ -575,6 +579,18 @@ def _bids_filter(value, parser):
         action='store_true',
         help='Disable head-motion correction and use the uncorrected data.',
     )
+    g_hmc.add_argument(
+        '--petref',
+        default='template',
+        choices=['template', 'twa', 'sum', 'first5min'],
+        help=(
+            "Strategy for generating the PET reference. 'template' uses the "
+            "motion correction template, while 'twa' computes a time-weighted "
+            "average, 'sum' produces a summed image of the motion-corrected "
+            "series, and 'first5min' averages the early (0-5 minute) portion "
+            'of the acquisition.'
+        ),
+    )
 
     g_seg = parser.add_argument_group('Segmentation options')
     g_seg.add_argument(
@@ -758,8 +774,9 @@ def parse_args(args=None, namespace=None):
     parser = _build_parser()
     opts = parser.parse_args(args, namespace)
 
-    if getattr(opts, 'pet2anat_robust', False) and opts.pet2anat_dof != 6:
-        parser.error('--pet2anat-robust requires --pet2anat-dof=6.')
+    # Validate DoF constraints for registration methods
+    if opts.pet2anat_method in ('robust', 'ants') and opts.pet2anat_dof != 6:
+        parser.error(f'--pet2anat-method {opts.pet2anat_method} requires --pet2anat-dof=6.')
 
     if opts.config_file:
         skip = {} if opts.reports_only else {'execution': ('run_uuid',)}

petprep/config.py

@@ -556,8 +556,10 @@ class workflow(_Config):
     """Degrees of freedom of the PET-to-anatomical registration steps."""
     pet2anat_init = 'auto'
     """Initial transform for PET-to-anatomical registration."""
-    pet2anat_robust = False
-    """Use ``mri_robust_register`` for PET-to-anatomical alignment."""
+    pet2anat_method: str = 'mri_coreg'
+    """PET-to-anatomical registration method (mri_coreg, robust, or ants)."""
+    petref: str = 'template'
+    """Strategy for building the PET reference (``'template'``, ``'twa'``, ``'sum'`` or ``'first5min'``)."""
     cifti_output = None
     """Generate HCP Grayordinates, accepts either ``'91k'`` (default) or ``'170k'``."""
     hires = None

petprep/interfaces/reports.py

@@ -98,6 +98,7 @@ def get_world_pedir(orientation: str, pe_dir: str) -> str:
 \t\t<ul class="elem-desc">
 \t\t\t<li>Original orientation: {ornt}</li>
 \t\t\t<li>Registration: {registration}</li>
+\t\t\t<li>Reference image: {reference}</li>
 \t\t\t<li>Time zero: {time_zero}</li>
 \t\t\t<li>Radiotracer: {radiotracer}</li>
 \t\t\t<li>Injected dose: {dose} {dose_units}</li>
@@ -228,6 +229,7 @@ class FunctionalSummaryInputSpec(TraitedSpec):
     registration = traits.Enum(
         'mri_coreg',
         'mri_robust_register',
+        'ants_registration',
         'Precomputed',
         mandatory=True,
         desc='PET/anatomical registration method',
@@ -237,6 +239,18 @@ class FunctionalSummaryInputSpec(TraitedSpec):
     )
     orientation = traits.Str(mandatory=True, desc='Orientation of the voxel axes')
     metadata = traits.Dict(desc='PET metadata dictionary')
+    petref_strategy = traits.Enum(
+        'template',
+        'twa',
+        'sum',
+        'first5min',
+        mandatory=True,
+        desc='PET reference generation strategy',
+    )
+    requested_petref_strategy = traits.Enum(
+        'template', 'twa', 'sum', 'first5min', desc='User-requested PET reference strategy'
+    )
+    hmc_disabled = traits.Bool(False, desc='Head motion correction disabled')
 
 
 class FunctionalSummary(SummaryInterface):
@@ -249,8 +263,25 @@ def _generate_segment(self):
             reg = 'Precomputed affine transformation'
         elif self.inputs.registration == 'mri_coreg':
             reg = f'FreeSurfer <code>mri_coreg</code> - {dof} dof'
+        elif self.inputs.registration == 'ants_registration':
+            reg = f'ANTs <code>ants_registration</code> ({dof} DoF)'
+        elif self.inputs.registration == 'mri_robust_register':
+            reg = 'FreeSurfer <code>mri_robust_register</code> (NMI cost)'
         else:
-            reg = 'FreeSurfer <code>mri_robust_register</code> (ROBENT cost)'
+            reg = f'Unknown registration method: {self.inputs.registration}'
+
+        reference_map = {
+            'template': 'Motion correction template',
+            'twa': 'Time-weighted average of motion-corrected series',
+            'sum': 'Summed motion-corrected series',
+            'first5min': 'Early (0-5 minute) average of motion-corrected series',
+        }
+        petref_strategy = reference_map.get(self.inputs.petref_strategy, 'Unknown')
+        requested = getattr(self.inputs, 'requested_petref_strategy', None)
+        if requested and requested != self.inputs.petref_strategy:
+            petref_strategy += f" (requested '{requested}')"
+        if self.inputs.hmc_disabled:
+            petref_strategy += ' (head motion correction disabled)'
 
         meta = self.inputs.metadata or {}
         time_zero = meta.get('TimeZero', None)
@@ -288,6 +319,7 @@ def _generate_segment(self):
 
         return FUNCTIONAL_TEMPLATE.format(
             registration=reg,
+            reference=petref_strategy,
             ornt=self.inputs.orientation,
             # Use the metadata dictionary to fill in the details
             time_zero=time_zero,

petprep/interfaces/tests/test_reports.py

@@ -76,7 +76,7 @@ def test_subject_summary_handles_missing_task(tmp_path):
 
 @pytest.mark.parametrize(
     'registration',
-    ['mri_coreg', 'mri_robust_register'],
+    ['mri_coreg', 'mri_robust_register', 'ants_registration'],
 )
 def test_functional_summary_with_metadata(registration):
     from ..reports import FunctionalSummary
@@ -85,6 +85,7 @@ def test_functional_summary_with_metadata(registration):
         registration=registration,
         registration_dof=6,
         orientation='RAS',
+        petref_strategy='template',
         metadata={
             'TracerName': 'DASB',
             'TracerRadionuclide': '[11C]',
@@ -97,6 +98,7 @@ def test_functional_summary_with_metadata(registration):
 
     segment = summary._generate_segment()
     assert registration in segment
+    assert 'Reference image: Motion correction template' in segment
     assert 'Radiotracer: [11C]DASB' in segment
     assert 'Injected dose: 100 MBq' in segment
     assert 'Number of frames: 2' in segment