FIX: Generate composite transform when using multi-step registration

Author: mgxd

nibabies/workflows/anatomical/fit.py

@@ -991,7 +991,6 @@ def init_infant_anat_fit_wf(
                 ('anat2std_xfm', 'inputnode.anat2std_xfm'),
                 ('std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
-            (anat_buffer, concat_reg_wf, [('anat_preproc', 'inputnode.anat_preproc')]),
             (sourcefile_buffer, ds_concat_reg_wf, [
                 ('anat_source_files', 'inputnode.source_files')
             ]),
@@ -1909,7 +1908,6 @@ def init_infant_single_anat_fit_wf(
                 ('anat2std_xfm', 'inputnode.anat2std_xfm'),
                 ('std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
-            (anat_buffer, concat_reg_wf, [('anat_preproc', 'inputnode.anat_preproc')]),
             (sourcefile_buffer, ds_concat_reg_wf, [
                 ('anat_source_files', 'inputnode.source_files')
             ]),

nibabies/workflows/anatomical/registration.py

@@ -15,15 +15,11 @@
 from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
 from smriprep.workflows.fit.registration import (
     TemplateDesc,
-    TemplateFlowSelect,
     _fmt_cohort,
     get_metadata,
     tf_ver,
 )
 
-from nibabies.config import DEFAULT_MEMORY_MIN_GB
-from nibabies.interfaces.patches import ConcatXFM
-
 
 def init_coregistration_wf(
     *,
@@ -312,7 +308,7 @@ def init_concat_registrations_wf(
     name='concat_registrations_wf',
 ):
     """
-    Concatenate two transforms to produce a single transform, from native to ``template``.
+    Concatenate two transforms to produce a single composite transform from native to template.
 
     Parameters
     ----------
@@ -347,6 +343,8 @@ def init_concat_registrations_wf(
         further use in downstream nodes.
 
     """
+    from nibabies.interfaces.patches import CompositeTransformUtil
+
     ntpls = len(templates)
     workflow = Workflow(name=name)
 
@@ -384,9 +382,7 @@ def init_concat_registrations_wf(
             workflow.__desc__ += '.\n' if template == templates[-1] else ', '
 
     inputnode = pe.Node(
-        niu.IdentityInterface(
-            fields=['template', 'anat_preproc', 'anat2std_xfm', 'intermediate', 'std2anat_xfm']
-        ),
+        niu.IdentityInterface(fields=['template', 'intermediate', 'anat2std_xfm', 'std2anat_xfm']),
         name='inputnode',
     )
     inputnode.inputs.template = templates
@@ -413,29 +409,39 @@ def init_concat_registrations_wf(
         TemplateDesc(), run_without_submitting=True, iterfield='template', name='split_desc'
     )
 
-    tf_select = pe.MapNode(
-        TemplateFlowSelect(resolution=1),
-        name='tf_select',
-        run_without_submitting=True,
-        iterfield=['template', 'template_spec'],
+    merge_anat2std = pe.Node(niu.Merge(2), name='merge_anat2std', run_without_submitting=True)
+    merge_std2anat = merge_anat2std.clone('merge_std2anat')
+
+    disassemble_anat2std = pe.MapNode(
+        CompositeTransformUtil(process='disassemble', output_prefix='anat2std'),
+        iterfield=['in_file'],
+        name='disassemble_anat2std',
     )
 
-    merge_anat2std = pe.MapNode(
-        niu.Merge(2), name='merge_anat2std', iterfield=['in1', 'in2'], run_without_submitting=True
+    disassemble_std2anat = pe.MapNode(
+        CompositeTransformUtil(process='disassemble', output_prefix='std2anat', inverse=True),
+        iterfield=['in_file'],
+        name='disassemble_std2anat',
+    )
+
+    order_anat2std_composites = pe.Node(
+        niu.Function(function=_order_composites),
+        name='order_anat2std_composites',
+    )
+
+    order_std2anat_composites = pe.Node(
+        niu.Function(function=_order_composites),
+        name='order_std2anat_composites',
     )
-    merge_std2anat = merge_anat2std.clone('merge_std2anat')
 
-    concat_anat2std = pe.MapNode(
-        ConcatXFM(),
-        name='concat_anat2std',
-        mem_gb=DEFAULT_MEMORY_MIN_GB,
-        iterfield=['transforms', 'reference_image'],
+    assemble_anat2std = pe.Node(
+        CompositeTransformUtil(process='assemble', out_file='anat2std.h5'),
+        name='assemble_anat2std',
     )
-    concat_std2anat = pe.MapNode(
-        ConcatXFM(),
-        name='concat_std2anat',
-        mem_gb=DEFAULT_MEMORY_MIN_GB,
-        iterfield=['transforms', 'reference_image'],
+
+    assemble_std2anat = pe.Node(
+        CompositeTransformUtil(process='assemble', out_file='std2anat.h5'),
+        name='assemble_std2anat',
     )
 
     fmt_cohort = pe.MapNode(
@@ -446,24 +452,30 @@ def init_concat_registrations_wf(
     )
 
     workflow.connect([
+        # Template concatenation
         (inputnode, merge_anat2std, [('anat2std_xfm', 'in2')]),
         (inputnode, merge_std2anat, [('std2anat_xfm', 'in2')]),
-        (inputnode, concat_std2anat, [('anat_preproc', 'reference_image')]),
         (inputnode, intermed_xfms, [('intermediate', 'intermediate')]),
         (inputnode, intermed_xfms, [('template', 'std')]),
-
         (intermed_xfms, merge_anat2std, [('int2std_xfm', 'in1')]),
         (intermed_xfms, merge_std2anat, [('std2int_xfm', 'in1')]),
+        (merge_anat2std, disassemble_anat2std, [('out', 'in_file')]),
+        (merge_std2anat, disassemble_std2anat, [('out', 'in_file')]),
+        (disassemble_anat2std, order_anat2std_composites, [
+            ('affine_transform', 'affines'),
+            ('displacement_field', 'displacements'),
+        ]),
+        (disassemble_std2anat, order_std2anat_composites, [
+            ('affine_transform', 'affines'),
+            ('displacement_field', 'displacements'),
+        ]),
+        (order_anat2std_composites, assemble_anat2std, [('out', 'in_file')]),
+        (order_std2anat_composites, assemble_std2anat, [('out', 'in_file')]),
+        (assemble_anat2std, outputnode, [('out_file', 'anat2std_xfm')]),
+        (assemble_std2anat, outputnode, [('out_file', 'std2anat_xfm')]),
 
-        (merge_anat2std, concat_anat2std, [('out', 'transforms')]),
-        (merge_std2anat, concat_std2anat, [('out', 'transforms')]),
-
+        # Template name wrangling
         (inputnode, split_desc, [('template', 'template')]),
-        (split_desc, tf_select, [
-            ('name', 'template'),
-            ('spec', 'template_spec'),
-        ]),
-        (tf_select, concat_anat2std, [('t1w_file', 'reference_image')]),
         (split_desc, fmt_cohort, [
             ('name', 'template'),
             ('spec', 'spec'),
@@ -472,8 +484,6 @@ def init_concat_registrations_wf(
             ('template', 'template'),
             ('spec', 'template_spec'),
         ]),
-        (concat_anat2std, outputnode, [('out_xfm', 'anat2std_xfm')]),
-        (concat_std2anat, outputnode, [('out_xfm', 'std2anat_xfm')]),
     ])  # fmt:skip
 
     return workflow
@@ -507,3 +517,7 @@ def _load_intermediate_xfms(intermediate, std):
     )
 
     return int2std, std2int
+
+
+def _order_composites(affines, displacements):
+    return [affines[0], displacements[0], affines[1], displacements[1]]