diff --git a/pytim/gitim.py b/pytim/gitim.py index 0b5da35d..76fc887a 100644 --- a/pytim/gitim.py +++ b/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 diff --git a/pytim/interface.py b/pytim/interface.py index 25e910a3..686a307b 100644 --- a/pytim/interface.py +++ b/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 _(): - >>> 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)) + + + >>> 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)) diff --git a/pytim/itim.py b/pytim/itim.py index 1e154d9a..66f3bdda 100644 --- a/pytim/itim.py +++ b/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 diff --git a/pytim/sasa.py b/pytim/sasa.py index 30c5feac..86e74dc7 100644 --- a/pytim/sasa.py +++ b/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: diff --git a/pytim/utilities_dbscan.py b/pytim/utilities_dbscan.py index 02ae0b14..ab06b174 100644 --- a/pytim/utilities_dbscan.py +++ b/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\ diff --git a/pytim/willard_chandler.py b/pytim/willard_chandler.py index 64dfc589..162de3c6 100644 --- a/pytim/willard_chandler.py +++ b/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)