add option to save transform

Author: salma1601

nipype/interfaces/afni/preprocess.py

@@ -2758,7 +2758,8 @@ class WarpInputSpec(AFNICommandInputSpec):
         name_template='%s_warp',
         desc='output image file name',
         argstr='-prefix %s',
-        name_source='in_file')
+        name_source='in_file',
+        keep_extension=True)
     tta2mni = traits.Bool(
         desc='transform dataset from Talairach to MNI152', argstr='-tta2mni')
     mni2tta = traits.Bool(
@@ -2789,6 +2790,13 @@ class WarpInputSpec(AFNICommandInputSpec):
         argstr='-zpad %d')
     verbose = traits.Bool(
         desc='Print out some information along the way.', argstr='-verb')
+    save_warp = traits.Bool(
+        False, usedefault=True, desc='save warp as .mat file')
+
+
+class WarpOutputSpec(TraitedSpec):
+    out_file = File(desc='Warped file.', exists=True)
+    warp_file = File(desc='warp transfrom .mat file')
 
 
 class Warp(AFNICommand):
@@ -2820,7 +2828,48 @@ class Warp(AFNICommand):
     """
     _cmd = '3dWarp'
     input_spec = WarpInputSpec
-    output_spec = AFNICommandOutputSpec
+    output_spec = WarpOutputSpec
+
+    def _run_interface(self, runtime):
+        runtime = super(Warp, self)._run_interface(runtime)
+
+        if self.inputs.save_warp:
+            import numpy as np
+            if not self.inputs.out_file:
+                warp_file = fname_presuffix(self.inputs.in_file,
+                                            suffix='_warp_transform.mat',
+                                            use_ext=False)                        
+            else:
+                warp_file = fname_presuffix(self.inputs.out_file,
+                                            suffix='_transform.mat',
+                                            use_ext=False)
+            np.savetxt(warp_file, [runtime.stdout], fmt=str('%s'))
+        return runtime
+
+    def _list_outputs(self):
+        outputs = self.output_spec().get()
+        if not self.inputs.out_file:
+            fname, ext = os.path.splitext(self.inputs.in_file)
+            if '.gz' in ext:
+                _, ext2 = os.path.splitext(fname)
+                ext = ext2 + ext
+            out_file = self._gen_fname(self.inputs.in_file, suffix='_warp',
+                                       ext=ext)
+            outputs['out_file'] = os.path.abspath(out_file)
+            if self.inputs.save_warp:
+                warp_file = fname_presuffix(self.inputs.in_file,
+                                            suffix='_warp_transform.mat',
+                                            use_ext=False)
+                outputs['warp_file'] = os.path.abspath(warp_file)
+        else:
+            outputs['out_file'] = os.path.abspath(self.inputs.out_file)
+            if self.inputs.save_warp:
+                warp_file = fname_presuffix(self.inputs.out_file,
+                                            suffix='_transform.mat',
+                                            use_ext=False)
+                outputs['warp_file'] = os.path.abspath(warp_file)
+
+        return outputs
 
 
 class QwarpPlusMinusInputSpec(CommandLineInputSpec):