From 1f6f6155d3d01be18b2ce8da41d7672caa31f164 Mon Sep 17 00:00:00 2001
From: Mathias Goncalves <goncalves.mathias@gmail.com>
Date: Wed, 11 Dec 2024 15:17:54 -0500
Subject: [PATCH] ENH: Support spatial normalization to alternative modalities

---
 smriprep/workflows/fit/registration.py | 50 ++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 3 deletions(-)

diff --git a/smriprep/workflows/fit/registration.py b/smriprep/workflows/fit/registration.py
index fc78575721..dfbefa1154 100644
--- a/smriprep/workflows/fit/registration.py
+++ b/smriprep/workflows/fit/registration.py
@@ -38,9 +38,10 @@
 
 def init_register_template_wf(
     *,
-    sloppy,
-    omp_nthreads,
-    templates,
+    sloppy: bool,
+    omp_nthreads: int,
+    templates: list[str],
+    image_type: str = 'T1w',
     name='register_template_wf',
 ):
     """
@@ -89,6 +90,8 @@ def init_register_template_wf(
         input domain to enable standardization of lesioned brains.
     template
         Template name and specification
+    image_type
+        Moving image modality
 
     Outputs
     -------
@@ -170,6 +173,15 @@ def init_register_template_wf(
         run_without_submitting=True,
     )
 
+    set_reference = pe.Node(
+        niu.Function(
+            function=_set_reference,
+            output_names=['reference_type', 'use_histogram_matching'],
+        ),
+        name='set_reference',
+    )
+    set_reference.inputs.image_type = image_type
+
     # With the improvements from nipreps/niworkflows#342 this truncation is now necessary
     trunc_mov = pe.Node(
         ants.ImageMath(operation='TruncateImageIntensity', op2='0.01 0.999 256'),
@@ -203,6 +215,14 @@ def init_register_template_wf(
             ('name', 'template'),
             ('spec', 'template_spec'),
         ]),
+        (tf_select, set_reference, [
+            ('t1w_file', 'template_t1w'),
+            ('t2w_file', 'template_t2w'),
+        ])
+        (set_reference, registration, [
+            ('reference_type', 'reference'),
+            ('use_histogram_matching', 'use_histogram_matching'),
+        ])
         (split_desc, registration, [
             ('name', 'template'),
             ('spec', 'template_spec'),
@@ -243,3 +263,27 @@ def _fmt_cohort(template, spec):
     if cohort is not None:
         template = f'{template}:cohort-{cohort}'
     return template, spec
+
+
+def _set_reference(image_type, template_t1w, template_t2w):
+    """
+    Determine the normalization reference and whether histogram matching will be used.
+
+    Parameters
+    ----------
+    image_type : MR image type of anatomical reference (T1w, T2w)
+    template_t1w : T1w file
+    template_t2w : T2w file or undefined
+
+    Returns
+    -------
+    reference_type : modality of template reference (T1w, T2w)
+    histogram_matching : False when image_type does not match reference image, otherwise Undefined
+    """
+    from nipype.interfaces.base import Undefined
+
+    if image_type == 'T2w':
+        if template_t2w:
+            return 'T2w', Undefined
+        return 'T1w', False
+    return 'T1w', Undefined