Merge pull request #2005 from mmodat/2004-niftyreg-base-intf

Fixing the base interface _run_interface function for NiftyRegCommand

Author: satra

nipype/interfaces/niftyreg/base.py

@@ -26,7 +26,7 @@
 import subprocess
 from warnings import warn
 
-from ..base import CommandLine, isdefined, CommandLineInputSpec, traits
+from ..base import CommandLine, CommandLineInputSpec, traits, Undefined
 from ...utils.filemanip import split_filename
 
 
@@ -47,8 +47,9 @@ def no_nifty_package(cmd='reg_f3d'):
 class NiftyRegCommandInputSpec(CommandLineInputSpec):
     """Input Spec for niftyreg interfaces."""
     # Set the number of omp thread to use
-    omp_core_val = traits.Int(desc='Number of openmp thread to use',
-                              argstr='-omp %i')
+    omp_core_val = traits.Int(int(os.environ.get('OMP_NUM_THREADS', '1')),
+                              desc='Number of openmp thread to use',
+                              argstr='-omp %i', usedefault=True)
 
 
 class NiftyRegCommand(CommandLine):
@@ -58,7 +59,10 @@ class NiftyRegCommand(CommandLine):
     _suffix = '_nr'
     _min_version = '1.5.30'
 
+    input_spec = NiftyRegCommandInputSpec
+
     def __init__(self, required_version=None, **inputs):
+        self.num_threads = 1
         super(NiftyRegCommand, self).__init__(**inputs)
         self.required_version = required_version
         _version = self.get_version()
@@ -73,6 +77,29 @@ def __init__(self, required_version=None, **inputs):
                     msg = 'The version of NiftyReg differs from the required'
                     msg += '(%s != %s)'
                     warn(msg % (_version, self.required_version))
+        self.inputs.on_trait_change(self._omp_update, 'omp_core_val')
+        self.inputs.on_trait_change(self._environ_update, 'environ')
+        self._omp_update()
+
+    def _omp_update(self):
+        if self.inputs.omp_core_val:
+            self.inputs.environ['OMP_NUM_THREADS'] = \
+                str(self.inputs.omp_core_val)
+            self.num_threads = self.inputs.omp_core_val
+        else:
+            if 'OMP_NUM_THREADS' in self.inputs.environ:
+                del self.inputs.environ['OMP_NUM_THREADS']
+            self.num_threads = 1
+
+    def _environ_update(self):
+        if self.inputs.environ:
+            if 'OMP_NUM_THREADS' in self.inputs.environ:
+                self.inputs.omp_core_val = \
+                    int(self.inputs.environ['OMP_NUM_THREADS'])
+            else:
+                self.inputs.omp_core_val = Undefined
+        else:
+            self.inputs.omp_core_val = Undefined
 
     def check_version(self):
         _version = self.get_version()
@@ -102,13 +129,6 @@ def version(self):
     def exists(self):
         return self.get_version() is not None
 
-    def _run_interface(self, runtime):
-        # Update num threads estimate from OMP_NUM_THREADS env var
-        # Default to 1 if not set
-        if not isdefined(self.inputs.environ['OMP_NUM_THREADS']):
-            self.inputs.environ['OMP_NUM_THREADS'] = self.num_threads
-        return super(NiftyRegCommand, self)._run_interface(runtime)
-
     def _format_arg(self, name, spec, value):
         if name == 'omp_core_val':
             self.numthreads = value

nipype/interfaces/niftyreg/tests/test_auto_NiftyRegCommand.py

@@ -12,6 +12,9 @@ def test_NiftyRegCommand_inputs():
     ignore_exception=dict(nohash=True,
     usedefault=True,
     ),
+    omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
+    ),
     terminal_output=dict(nohash=True,
     ),
     )

nipype/interfaces/niftyreg/tests/test_auto_RegAladin.py

@@ -46,6 +46,7 @@ def test_RegAladin_inputs():
     nosym_flag=dict(argstr='-noSym',
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     platform_val=dict(argstr='-platf %i',
     ),

nipype/interfaces/niftyreg/tests/test_auto_RegAverage.py

@@ -43,6 +43,7 @@ def test_RegAverage_inputs():
     usedefault=True,
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='%s',
     genfile=True,

nipype/interfaces/niftyreg/tests/test_auto_RegF3D.py

@@ -76,6 +76,7 @@ def test_RegF3D_inputs():
     noz_flag=dict(argstr='-noz',
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     pad_val=dict(argstr='-pad %f',
     ),

nipype/interfaces/niftyreg/tests/test_auto_RegJacobian.py

@@ -13,6 +13,7 @@ def test_RegJacobian_inputs():
     usedefault=True,
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='%s',
     name_source=['trans_file'],

nipype/interfaces/niftyreg/tests/test_auto_RegMeasure.py

@@ -19,6 +19,7 @@ def test_RegMeasure_inputs():
     mandatory=True,
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='-out %s',
     name_source=['flo_file'],

nipype/interfaces/niftyreg/tests/test_auto_RegResample.py

@@ -18,6 +18,7 @@ def test_RegResample_inputs():
     inter_val=dict(argstr='-inter %d',
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='%s',
     name_source=['flo_file'],

nipype/interfaces/niftyreg/tests/test_auto_RegTools.py

@@ -34,6 +34,7 @@ def test_RegTools_inputs():
     noscl_flag=dict(argstr='-noscl',
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='-out %s',
     name_source=['in_file'],

nipype/interfaces/niftyreg/tests/test_auto_RegTransform.py

@@ -57,6 +57,7 @@ def test_RegTransform_inputs():
     xor=['def_input', 'disp_input', 'flow_input', 'comp_input', 'upd_s_form_input', 'inv_aff_input', 'inv_nrr_input', 'half_input', 'aff_2_rig_input', 'flirt_2_nr_input'],
     ),
     omp_core_val=dict(argstr='-omp %i',
+    usedefault=True,
     ),
     out_file=dict(argstr='%s',
     genfile=True,