add MKL thread number controling class

Author: chrisgorgo

nipype/interfaces/base.py

@@ -10,6 +10,8 @@
 """
 from __future__ import print_function, division, unicode_literals, absolute_import
 from future import standard_library
+from py.path import ImportMismatchError
+
 standard_library.install_aliases()
 from builtins import range, object, open, str, bytes
 
@@ -1217,6 +1219,29 @@ def save_inputs_to_json(self, json_file):
             json.dump(inputs, fhandle, indent=4, ensure_ascii=False)
 
 
+class MKLInterfaceInputSpec(BaseInterfaceInputSpec):
+    mkl_num_threads = traits.Int(1, usedafult=True, nohash=True,
+                                 desc="Number of threads to use if mkl-service "
+                                      "package is installed.")
+
+
+class MKLInterface(BaseInterface):
+    """Interfaces that control number of threads available via MKL.
+    """
+
+    def _run_interface(self, runtime):
+        try:
+            import mkl
+        except ImportError:
+            iflogger.warning("mkl-service package not found - cannot control "
+                             "number of threads")
+            return runtime
+        else:
+            mkl.set_num_threads(self.inputs.mkl_num_threads)
+            self.num_threads = self.inputs.mkl_num_threads
+            return runtime
+
+
 class Stream(object):
     """Function to capture stdout and stderr streams with timestamps
 

nipype/interfaces/nipy/preprocess.py

@@ -21,7 +21,9 @@
 from ...utils.filemanip import split_filename, fname_presuffix
 from ..base import (TraitedSpec, BaseInterface, traits,
                     BaseInterfaceInputSpec, isdefined, File,
-                    InputMultiPath, OutputMultiPath)
+                    InputMultiPath, OutputMultiPath, MKLInterface,
+                    MKLInterfaceInputSpec
+                    )
 
 
 have_nipy = True
@@ -207,7 +209,7 @@ def _list_outputs(self):
         return outputs
 
 
-class SpaceTimeRealignerInputSpec(BaseInterfaceInputSpec):
+class SpaceTimeRealignerInputSpec(MKLInterfaceInputSpec):
 
     in_file = InputMultiPath(File(exists=True),
                              mandatory=True, min_ver='0.4.0.dev',
@@ -246,7 +248,7 @@ class SpaceTimeRealignerOutputSpec(TraitedSpec):
                                      "euler angles"))
 
 
-class SpaceTimeRealigner(BaseInterface):
+class SpaceTimeRealigner(MKLInterface):
     """Simultaneous motion and slice timing correction algorithm
 
     If slice_times is not specified, this algorithm performs spatial motion
@@ -290,6 +292,7 @@ def version(self):
         return nipy_version
 
     def _run_interface(self, runtime):
+        runtime = super(SpaceTimeRealigner, self)._run_interface()
         all_ims = [load_image(fname) for fname in self.inputs.in_file]
 
         if not isdefined(self.inputs.slice_times):