RF: Perform transform ordering when disassembling

Author: mgxd

nibabies/interfaces/patches.py

@@ -10,6 +10,9 @@
 from nipype.interfaces.ants.registration import (
     CompositeTransformUtilInputSpec as _CompositeTransformUtilInputSpec,
 )
+from nipype.interfaces.ants.registration import (
+    CompositeTransformUtilOutputSpec as _CompositeTransformUtilOutputSpec,
+)
 from nipype.interfaces.base import File, InputMultiObject, TraitedSpec, traits
 
 
@@ -114,35 +117,69 @@ def _list_outputs(self):
 
 
 class CompositeTransformUtilInputSpec(_CompositeTransformUtilInputSpec):
-    inverse = traits.Bool(
-        False,
+    order_transforms = traits.Bool(
+        True,
         usedefault=True,
-        desc='When disassembling an inverse component transform, the indexing will be reversed.',
+        desc='Order disassembled transforms into [Affine, Displacement] pairs.',
     )
 
 
+class CompositeTransformUtilOutputSpec(_CompositeTransformUtilOutputSpec):
+    out_transforms = traits.List(desc='list of transform components')
+
+
 class CompositeTransformUtil(_CompositeTransformUtil):
     """Outputs have changed in newer versions of ANTs."""
 
     input_spec = CompositeTransformUtilInputSpec
+    output_spec = CompositeTransformUtilOutputSpec
 
     def _list_outputs(self):
         outputs = self.output_spec().get()
 
-        # Index may change depending on forward/inverse transform
+        # Ordering may change depending on forward/inverse transform
         # Forward: <prefix>_00_AffineTransform.mat, <prefix>_01_DisplacementFieldTransform.nii.gz
         # Inverse: <prefix>_01_AffineTransform.mat, <prefix>_00_DisplacementFieldTransform.nii.gz
-        idx = ['00', '01']
-        if self.inputs.inverse:
-            idx = idx[::-1]
-
         if self.inputs.process == 'disassemble':
-            outputs['affine_transform'] = Path(
-                f'{self.inputs.output_prefix}_{idx[0]}_AffineTransform.mat'
-            ).absolute()
-            outputs['displacement_field'] = Path(
-                f'{self.inputs.output_prefix}_{idx[1]}_DisplacementFieldTransform.nii.gz'
-            ).absolute()
+            transforms = [
+                str(Path(x).absolute())
+                for x in sorted(Path().glob(f'{self.inputs.output_prefix}_*'))
+            ]
+
+            if self.inputs.order_transforms:
+                transforms = _order_xfms(transforms)
+            outputs['out_transforms'] = transforms
+
+            # Potentially could be more than one affine / displacement per composite transform...
+            outputs['affine_transform'] = [
+                x for x in transforms if 'AffineTransform' in Path(x).name
+            ][0]
+            outputs['displacement_field'] = [
+                x for x in transforms if 'DisplacementFieldTransform' in Path(x).name
+            ][0]
         elif self.inputs.process == 'assemble':
             outputs['out_file'] = Path(self.inputs.out_file).absolute()
         return outputs
+
+
+def _order_xfms(vals):
+    """
+    Assumes [affine, displacement] or [displacement, affine] transform pairs.
+
+    >>> _order_xfms(['DisplacementFieldTransform.nii.gz', 'AffineTransform.mat'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz']
+
+    >>> _order_xfms(['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz']
+
+    >>> _order_xfms(['DisplacementFieldTransform.nii.gz', 'AffineTransform.mat', \
+        'AffineTransform.mat'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz', 'AffineTransform.mat']
+    """
+    for i in range(0, len(vals) - 1, 2):
+        if (
+            'DisplacementFieldTransform' in Path(vals[i]).name
+            and 'AffineTransform' in Path(vals[i + 1]).name
+        ):
+            vals[i], vals[i + 1] = vals[i + 1], vals[i]
+    return vals

nibabies/workflows/anatomical/registration.py

@@ -419,26 +419,24 @@ def init_concat_registrations_wf(
     )
 
     disassemble_std2anat = pe.MapNode(
-        CompositeTransformUtil(process='disassemble', output_prefix='std2anat', inverse=True),
+        CompositeTransformUtil(process='disassemble', output_prefix='std2anat'),
         iterfield=['in_file'],
         name='disassemble_std2anat',
     )
 
-    order_anat2std_composites = pe.Node(
-        niu.Function(function=_order_composites),
-        name='order_anat2std_composites',
+    merge_anat2std_composites = pe.Node(
+        niu.Merge(1, ravel_inputs=True),
+        name='merge_anat2std_composites',
     )
-
-    order_std2anat_composites = pe.Node(
-        niu.Function(function=_order_composites),
-        name='order_std2anat_composites',
+    merge_std2anat_composites = pe.Node(
+        niu.Merge(1, ravel_inputs=True),
+        name='merge_std2anat_composites',
     )
 
     assemble_anat2std = pe.Node(
         CompositeTransformUtil(process='assemble', out_file='anat2std.h5'),
         name='assemble_anat2std',
     )
-
     assemble_std2anat = pe.Node(
         CompositeTransformUtil(process='assemble', out_file='std2anat.h5'),
         name='assemble_std2anat',
@@ -461,16 +459,10 @@ def init_concat_registrations_wf(
         (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')]),
+        (disassemble_anat2std, merge_anat2std_composites, [('out_transforms', 'in1')]),
+        (disassemble_std2anat, merge_std2anat_composites, [('out_transforms', 'in1')]),
+        (merge_anat2std_composites, assemble_anat2std, [('out', 'in_file')]),
+        (merge_std2anat_composites, assemble_std2anat, [('out', 'in_file')]),
         (assemble_anat2std, outputnode, [('out_file', 'anat2std_xfm')]),
         (assemble_std2anat, outputnode, [('out_file', 'std2anat_xfm')]),
 
@@ -517,7 +509,3 @@ def _load_intermediate_xfms(intermediate, std):
     )
 
     return int2std, std2int
-
-
-def _order_composites(affines, displacements):
-    return [affines[0], displacements[0], affines[1], displacements[1]]