ALL: allow to change number of threads used for kdtree

pytim/gitim.py

@@ -66,6 +66,8 @@ class GITIM(Interface):
         :param bool warnings:       Print warnings
         :param bool autoassign:     If true (default) detect the interface
                                     every time a new frame is selected.
+        :param int workers:         Passes the workers option to scipy.spatial.cKDTree.query_ball
+                                    If -1 is given all CPU threads are used (default: -1)
 
         Example:
 
@@ -127,13 +129,14 @@ def __init__(self,
                  warnings=False,
                  autoassign=True,
                  _noextrapoints=False,
+                 workers=-1,
                  **kargs):
 
         # this is just for debugging/testing
         self._noextrapoints = _noextrapoints
         self.autoassign = autoassign
         self.system = platform.system()
-
+        self.workers = workers
         self.do_center = centered
 
         self.biggest_cluster_only = biggest_cluster_only

pytim/interface.py

@@ -148,6 +148,11 @@ def _assign_symmetry(self, symmetry):
     def _define_cluster_group(self):
         self.universe.atoms.pack_into_box()
         self.cluster_group = self.universe.atoms[:0]  # empty
+        try: # for classes that are not yet implementing this, play safe
+            self.workers
+        except AttributeError:
+            self.workers=1
+
         if (self.cluster_cut is not None):
             cluster_cut = float(self.cluster_cut[0])
             # we start by adding the atoms in the smaller clusters
@@ -156,7 +161,7 @@ def _define_cluster_group(self):
                 for extra in self.extra_cluster_groups:
                     x_labels, x_counts, _ = utilities.do_cluster_analysis_dbscan(
                         extra, cluster_cut, self.cluster_threshold_density,
-                        self.molecular)
+                        self.molecular,workers=self.workers)
                     x_labels = np.array(x_labels)
                     x_label_max = np.argmax(x_counts)
                     x_ids_other = np.where(x_labels != x_label_max)[0]
@@ -171,7 +176,7 @@ def _define_cluster_group(self):
             # the smaller clusters of the other phase
             labels, counts, neighbors = utilities.do_cluster_analysis_dbscan(
                 self.cluster_group, cluster_cut,
-                self.cluster_threshold_density, self.molecular)
+                self.cluster_threshold_density, self.molecular,workers=self.workers)
             labels = np.array(labels)
 
             # counts is not necessarily ordered by size of cluster.
@@ -519,7 +524,11 @@ def _():
         <AtomGroup with 2025 atoms>
 
 
-        >>> inter = pytim.SASA( g, alpha=2.5, max_layers=2, cluster_cut=3.5, biggest_cluster_only=True, molecular=True)
+        >>> inter = pytim.SASA( g, alpha=2.5, max_layers=2, cluster_cut=3.5, biggest_cluster_only=True, molecular=True,workers=-1)
+        >>> print(repr(inter.atoms))
+        <AtomGroup with 855 atoms>
+
+        >>> inter = pytim.SASA( g, alpha=2.5, max_layers=2, cluster_cut=3.5, biggest_cluster_only=True, molecular=True,workers=1)
         >>> print(repr(inter.atoms))
         <AtomGroup with 855 atoms>
 

pytim/itim.py

@@ -63,6 +63,8 @@ class ITIM(Interface):
                                   (default 0.4 Angstrom)
         :param bool autoassign:   If true (default) detect the interface
                                   every time a new frame is selected.
+        :param int workers:       Passes the workers option to scipy.spatial.cKDTree.query_ball
+                                  If -1 is given all CPU threads are used (default: -1)
 
         Example:
 
@@ -193,9 +195,11 @@ def __init__(self,
                  warnings=False,
                  mesh=0.4,
                  autoassign=True,
+                 workers=-1,
                  **kargs):
 
         self.autoassign = autoassign
+        self.workers = workers
         self.biggest_cluster_only = True # necessary for ITIM
         self.symmetry = 'planar'
         self.do_center = centered

pytim/sasa.py

@@ -60,6 +60,8 @@ class SASA(GITIM):
         :param bool warnings:       Print warnings
         :param bool autoassign:     If true (default) detect the interface
                                     every time a new frame is selected.
+        :param int workers:         Passes the workers option to scipy.spatial.cKDTree.query_ball
+                                    If -1 is given all CPU threads are used (default: -1)
 
         Example:
 

pytim/utilities_dbscan.py

@@ -31,7 +31,7 @@ def determine_samples(threshold_density, cluster_cut, n_neighbors):
 def do_cluster_analysis_dbscan(group,
                                cluster_cut,
                                threshold_density=None,
-                               molecular=True):
+                               molecular=True,workers=1):
     """ Performs a cluster analysis using DBSCAN
 
         :returns [labels,counts,neighbors]: lists of the id of the cluster to
@@ -46,6 +46,9 @@ def do_cluster_analysis_dbscan(group,
         clusters. This is on average O(N log N) thanks to the O(log N)
         scaling of the kdtree.
 
+        To avoid problems in environments with limited resources, 
+        pass n_jobs=1 to start only one thread
+
     """
     box = group.universe.dimensions[:3]
 
@@ -56,7 +59,7 @@ def do_cluster_analysis_dbscan(group,
 
     neighborhoods = np.array([
         np.array(neighbors)
-        for neighbors in tree.query_ball_point(points, cluster_cut, n_jobs=-1)
+        for neighbors in tree.query_ball_point(points, cluster_cut, workers=workers)
     ])
     if len(neighborhoods.shape) != 1:
         raise ValueError("Error in do_cluster_analysis_DBSCAN(), the cutoff\

pytim/willard_chandler.py

@@ -55,6 +55,10 @@ class WillardChandler(Interface):
                                   mixed interfaces
         :param bool centered:     Center the  :py:obj:`group`
         :param bool warnings:     Print warnings
+        :param bool autoassign:   If true (default) detect the interface
+                                  every time a new frame is selected.
+        :param int workers:       Passes the workers option to scipy.spatial.cKDTree.query_ball
+                                  If -1 is given all CPU threads are used (default: -1)
 
         Example:
 
@@ -109,9 +113,11 @@ def __init__(self,
                  centered=False,
                  warnings=False,
                  autoassign=True,
+                 workers=-1,
                  **kargs):
 
         self.autoassign, self.do_center = autoassign, centered
+        self.workers = workers
         sanity = SanityCheck(self, warnings=warnings)
         sanity.assign_universe(universe, group)
         sanity.assign_alpha(alpha)