fix for issue #74 supplying n_jobs as a parameter for robust single linkage.

Author: lmcinnes

hdbscan/hdbscan_.py

@@ -613,7 +613,8 @@ class HDBSCAN (BaseEstimator, ClusterMixin):
 
     core_dist_n_jobs : int, optional
         Number of parallel jobs to run in core distance computations (if
-        supported by the specific algorithm).
+        supported by the specific algorithm). For ``core_dist_n_jobs``
+        below -1, (n_cpus + 1 + core_dist_n_jobs) are used.
         (default 4)
 
     allow_single_cluster : boolean

hdbscan/robust_single_linkage_.py

@@ -87,7 +87,11 @@ def _rsl_prims_balltree(X, k=5, alpha=1.4142135623730951, metric='euclidean', **
 
 
 def _rsl_boruvka_kdtree(X, k=5, alpha=1.0,
-                        metric='euclidean', leaf_size=40, **kwargs):
+                        metric='euclidean', leaf_size=40,
+                        core_dist_n_jobs=4, **kwargs):
+
+    if core_dist_n_jobs < 1:
+        core_dist_n_jobs = max(cpu_count() + 1 + core_dist_n_jobs, 1)
 
     dim = X.shape[0]
     min_samples = min(dim - 1, k)
@@ -104,9 +108,13 @@ def _rsl_boruvka_kdtree(X, k=5, alpha=1.0,
 
 
 def _rsl_boruvka_balltree(X, k=5, alpha=1.0,
-                          metric='euclidean', leaf_size=40, **kwargs):
+                          metric='euclidean', leaf_size=40,
+                          core_dist_n_jobs=4, **kwargs):
 
-    dim = X.shape[0]
+    if core_dist_n_jobs < 1:
+        core_dist_n_jobs = max(cpu_count() + 1 + core_dist_n_jobs, 1)
+
+     dim = X.shape[0]
     min_samples = min(dim - 1, k)
 
     tree = BallTree(X, metric=metric, leaf_size=leaf_size, **kwargs)
@@ -122,7 +130,8 @@ def _rsl_boruvka_balltree(X, k=5, alpha=1.0,
 
 def robust_single_linkage(X, cut, k=5, alpha=1.4142135623730951,
                           gamma=5, metric='euclidean', algorithm='best',
-                          memory=Memory(cachedir=None, verbose=0), leaf_size=40, **kwargs):
+                          memory=Memory(cachedir=None, verbose=0), leaf_size=40,
+                          core_dist_n_jobs=4, **kwargs):
     """Perform robust single linkage clustering from a vector array
     or distance matrix.
 
@@ -180,6 +189,12 @@ def robust_single_linkage(X, cut, k=5, alpha=1.4142135623730951,
         Leaf size for trees responsible for fast nearest
         neighbour queries. (default 40)
 
+    core_dist_n_jobs : int, optional
+        Number of parallel jobs to run in core distance computations (if
+        supported by the specific algorithm). For ``core_dist_n_jobs``
+        below -1, (n_cpus + 1 + core_dist_n_jobs) are used.
+        (default 4)
+
     Returns
     -------
     labels : array [n_samples]
@@ -232,10 +247,12 @@ def robust_single_linkage(X, cut, k=5, alpha=1.4142135623730951,
                 memory.cache(_rsl_prims_balltree)(X, k, alpha, metric, **kwargs)
         elif algorithm == 'boruvka_kdtree':
             single_linkage_tree = \
-                memory.cache(_rsl_boruvka_kdtree)(X, k, alpha, metric, leaf_size, **kwargs)
+                memory.cache(_rsl_boruvka_kdtree)(X, k, alpha, metric, leaf_size,
+                                                  core_dist_n_jobs, **kwargs)
         elif algorithm == 'boruvka_balltree':
             single_linkage_tree = \
-                memory.cache(_rsl_boruvka_balltree)(X, k, alpha, metric, leaf_size, **kwargs)
+                memory.cache(_rsl_boruvka_balltree)(X, k, alpha, metric, leaf_size,
+                                                    core_dist_n_jobs, **kwargs)
         else:
             raise TypeError('Unknown algorithm type %s specified' % algorithm)
     else:
@@ -249,21 +266,26 @@ def robust_single_linkage(X, cut, k=5, alpha=1.4142135623730951,
                     memory.cache(_rsl_prims_kdtree)(X, k, alpha, metric, **kwargs)
             else:
                 single_linkage_tree = \
-                    memory.cache(_rsl_boruvka_kdtree)(X, k, alpha, metric, **kwargs)
+                    memory.cache(_rsl_boruvka_kdtree)(X, k, alpha, metric,
+                                                        leaf_size,
+                                                        core_dist_n_jobs,
+                                                        **kwargs)
         else:  # Metric is a valid BallTree metric
             # Need heuristic to decide when to go to boruvka; still debugging for now
             if X.shape[1] > 128:
                 single_linkage_tree = \
                     memory.cache(_rsl_prims_kdtree)(X, k, alpha, metric, **kwargs)
             else:
                 single_linkage_tree = \
-                    memory.cache(_rsl_boruvka_balltree)(X, k, alpha, metric, **kwargs)
+                    memory.cache(_rsl_boruvka_balltree)(X, k, alpha, metric,
+                                                        leaf_size,
+                                                        core_dist_n_jobs,
+                                                        **kwargs)
 
     labels = single_linkage_tree.get_clusters(cut, gamma)
 
     return labels, single_linkage_tree
 
-
 class RobustSingleLinkage(BaseEstimator, ClusterMixin):
     """Perform robust single linkage clustering from a vector array
     or distance matrix.
@@ -317,6 +339,13 @@ class RobustSingleLinkage(BaseEstimator, ClusterMixin):
             * ``large_kdtree``
             * ``large_kdtree_fastcluster``
 
+
+    core_dist_n_jobs : int, optional
+        Number of parallel jobs to run in core distance computations (if
+        supported by the specific algorithm). For ``core_dist_n_jobs``
+        below -1, (n_cpus + 1 + core_dist_n_jobs) are used.
+        (default 4)
+
     Attributes
     -------
     labels_ : array [n_samples]
@@ -339,14 +368,15 @@ class RobustSingleLinkage(BaseEstimator, ClusterMixin):
     """
 
     def __init__(self, cut=0.4, k=5, alpha=1.4142135623730951, gamma=5, metric='euclidean',
-                 algorithm='best', **kwargs):
+                 algorithm='best', core_dist_n_jobs=4, **kwargs):
 
         self.cut = cut
         self.k = k
         self.alpha = alpha
         self.gamma = gamma
         self.metric = metric
         self.algorithm = algorithm
+        self.core_dist_n_jobs = core_dist_n_jobs
 
         self._metric_kwargs = kwargs